{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Studio notebook lifecycle configurations\n",
    "\n",
    "This notebook implements a method for configuring dependencies within SageMaker Studio by using Studio Lifecycle Configuration whereby dependencies are installed each time a kernel starts-up.\n",
    "\n",
    "For this notebook, please use the <b>Data Science</b> image and <b>Python 3</b> kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    This notebook requires the <b>Studio Execution Role</b> to have <b>sagemaker:UpdateDomain</b> and <b>sagemaker:UpdateUserProfile</b> permissions\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For `UpdateDomain` permission you can use the following inline policy for the execution role:\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"sagemaker:UpdateDomain\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:sagemaker:us-east-1:<ACCOUNT_ID>:domain/<DOMAIN_ID>\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "A sample for the `UpdateUserProfile` policy:\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"sagemaker:UpdateUserProfile\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:sagemaker:us-east-1:<ACCOUNT_ID>:user-profile/<DOMAIN_ID>/*\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "import base64\n",
    "import boto3\n",
    "import botocore.exceptions\n",
    "import sagemaker\n",
    "\n",
    "client = boto3.client('sagemaker')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Studio Lifecyle Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "StudioLifecycleConfigName = 'InstallDependencies-%004x' % random.getrandbits(16)\n",
    "\n",
    "StudioLifecycleConfigContent = \"\"\"# This script installs a single pip package on a SageMaker Studio Kernel Application\n",
    "#!/bin/bash\n",
    "set -eux\n",
    "# PARAMETERS\n",
    "PACKAGE=pyarrow\n",
    "pip install --upgrade $PACKAGE\"\"\"\n",
    "\n",
    "message_bytes = StudioLifecycleConfigContent.encode('ascii')\n",
    "base64_bytes = base64.b64encode(message_bytes)\n",
    "\n",
    "studio_lifecycle_config = client.create_studio_lifecycle_config(\n",
    "    StudioLifecycleConfigName=StudioLifecycleConfigName,\n",
    "    StudioLifecycleConfigContent=base64_bytes.decode('ascii'),\n",
    "    StudioLifecycleConfigAppType='KernelGateway',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fetch the SageMaker Studio domain ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "NOTEBOOK_METADATA_FILE = \"/opt/ml/metadata/resource-metadata.json\"\n",
    "domain_id = None\n",
    "\n",
    "if os.path.exists(NOTEBOOK_METADATA_FILE):\n",
    "    with open(NOTEBOOK_METADATA_FILE, \"rb\") as f:\n",
    "        domain_id = json.loads(f.read()).get('DomainId')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There two possible options for the next step, option 1 is adding the lifecycle configuration to the SageMaker domain so that all profiles within that domain inherit it and option 2 is to add the lifecycle configuration to a user profile so that it is scoped to that specific profile."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Option 1: Add Studio Lifecyle Configuration to the SageMaker domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "domain = client.describe_domain(\n",
    "    DomainId=domain_id\n",
    ")\n",
    "\n",
    "try:\n",
    "    lifecycle_config_arns = domain['DefaultUserSettings']['KernelGatewayAppSettings']['LifecycleConfigArns']\n",
    "except KeyError as e:\n",
    "    lifecycle_config_arns = []\n",
    "\n",
    "lifecycle_config_arns.append(studio_lifecycle_config['StudioLifecycleConfigArn'])\n",
    "\n",
    "try:\n",
    "    client.update_domain(\n",
    "        DomainId=domain_id,\n",
    "        DefaultUserSettings={\n",
    "            'KernelGatewayAppSettings': {\n",
    "                'LifecycleConfigArns': lifecycle_config_arns,\n",
    "            }\n",
    "        }\n",
    "    )\n",
    "except botocore.exceptions.ClientError as e:\n",
    "    role = sagemaker.get_execution_role()\n",
    "    print(f\"SageMaker Studio execution role ({role}) does not have permission to perform UpdateDomain operation. Please add 'sagemaker:UpdateDomain' permission to the role.\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'LifecycleConfigArns': ['arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies3',\n",
       "  'arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies4',\n",
       "  'arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies-6566',\n",
       "  'arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies-c9d6']}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List domain-wide lifecycle configurations\n",
    "domain = client.describe_domain(\n",
    "    DomainId=domain_id\n",
    ")\n",
    "domain['DefaultUserSettings']['KernelGatewayAppSettings']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: Add Studio Lifecyle Configuration to a specific user profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'DomainId': 'd-hrcfizeddzyj',\n",
       "  'UserProfileName': 'domain-3-user-1',\n",
       "  'Status': 'InService',\n",
       "  'CreationTime': datetime.datetime(2022, 12, 11, 17, 58, 11, 59000, tzinfo=tzlocal()),\n",
       "  'LastModifiedTime': datetime.datetime(2023, 2, 1, 10, 42, 54, 569000, tzinfo=tzlocal())}]"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List all user profiles in the current domain\n",
    "client.list_user_profiles(DomainIdEquals=domain_id)['UserProfiles']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Select user profile from the list above\n",
    "user_profile_name = 'domain-3-user-1'\n",
    "\n",
    "\n",
    "user_profile = client.describe_user_profile(\n",
    "    DomainId=domain_id,\n",
    "    UserProfileName=user_profile_name\n",
    ")\n",
    "\n",
    "try:\n",
    "    lifecycle_config_arns = user_profile['UserSettings']['KernelGatewayAppSettings']['LifecycleConfigArns']\n",
    "except KeyError as e:\n",
    "    lifecycle_config_arns = []\n",
    "lifecycle_config_arns.append(studio_lifecycle_config['StudioLifecycleConfigArn'])\n",
    "\n",
    "try:\n",
    "    client.update_user_profile(\n",
    "        DomainId=domain_id,\n",
    "        UserProfileName=user_profile_name,\n",
    "        UserSettings={\n",
    "            'KernelGatewayAppSettings': {\n",
    "                'LifecycleConfigArns': lifecycle_config_arns,\n",
    "            }\n",
    "        }\n",
    "    )\n",
    "except botocore.exceptions.ClientError as e:\n",
    "    role = sagemaker.get_execution_role()\n",
    "    print(f\"SageMaker Studio execution role ({role}) does not have permission to perform UpdateUserProfile operation. Please add 'sagemaker:UpdateUserProfile' permission to the role.\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies-c9d6',\n",
       " 'arn:aws:sagemaker:us-east-1:906545278380:studio-lifecycle-config/installdependencies-3ad2']"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List user profile-specific lifecycle configuration\n",
    "user_profile = client.describe_user_profile(\n",
    "    DomainId=domain_id,\n",
    "    UserProfileName=user_profile_name\n",
    ")\n",
    "\n",
    "user_profile['UserSettings']['KernelGatewayAppSettings']['LifecycleConfigArns']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Test lifecycle configurations\n",
    "Now that all the configuration is done, it is time to test the script within Studio. To do this, launch the Studio and on the Launcher tab, locate the Notebooks and compute resources section and click on the Change environment to select the lifecycle configuration you created.\n",
    "\n",
    "![](../img/studio-create-notebook.png)\n",
    "\n",
    "On the Change environment popup, change the Start-up script option to be the lifecycle configuration created in previous steps.\n",
    "\n",
    "![](../img/change-environment-2.png)\n",
    "\n",
    "Back in the Launcher window, click on the Create notebook option.\n",
    "You can also set the Lifecycle configuration to be run by default in the Lifecycle configurations for personal Studio apps section of the Domain page.\n",
    "Within the new notebook, the dependencies installed in the start-up script will be available.\n",
    "\n",
    "![](../img/lifecycle-config-working.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resources\n",
    "Check out the following links for more information:\n",
    "- [Use Lifecycle Configurations with Amazon SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-lcc.html)\n",
    "- [Studio Lifecycle Configuration examples GitHub repository](https://github.com/aws-samples/sagemaker-studio-lifecycle-config-examples)\n",
    "- [Customize Amazon SageMaker Studio using Lifecycle Configurations](http://aws.amazon.com/blogs/machine-learning/customize-amazon-sagemaker-studio-using-lifecycle-configurations/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Shutdown kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<p><b>Shutting down your kernel for this notebook to release resources.</b></p>\n",
       "<button class=\"sm-command-button\" data-commandlinker-command=\"kernelmenu:shutdown\" style=\"display:none;\">Shutdown Kernel</button>\n",
       "        \n",
       "<script>\n",
       "try {\n",
       "    els = document.getElementsByClassName(\"sm-command-button\");\n",
       "    els[0].click();\n",
       "}\n",
       "catch(err) {\n",
       "    // NoOp\n",
       "}    \n",
       "</script>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "\n",
    "<p><b>Shutting down your kernel for this notebook to release resources.</b></p>\n",
    "<button class=\"sm-command-button\" data-commandlinker-command=\"kernelmenu:shutdown\" style=\"display:none;\">Shutdown Kernel</button>\n",
    "        \n",
    "<script>\n",
    "try {\n",
    "    els = document.getElementsByClassName(\"sm-command-button\");\n",
    "    els[0].click();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}    \n",
    "</script>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}