   ml_ops/sm-mlflow_setup/sm-mlflow_setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "02127090-ee33-4005-b5af-5f4e386ed1a6",
   "metadata": {},
   "source": [
    "# How to Setup Amazon SageMaker with MLflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "557f10ee-714a-4378-9493-abe2cd010754",
   "metadata": {},
   "source": [
    "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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f09f362-71a7-409c-a4c7-0ee5e59c1581",
   "metadata": {},
   "source": [
    "## Updates and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd83cc42-fc1e-49cd-88e3-7a685add2404",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U --quiet boto3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f932a722-a2cd-4aca-bdc0-d00553439966",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "329cf85a-352c-4f55-8e2a-4771a26fbe70",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sagemaker\n",
    "import boto3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95c2869e-1845-4534-bf97-d530b5c27c48",
   "metadata": {},
   "source": [
    "Session variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece22344-f747-4fb2-9051-3640dd95dd6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = sagemaker.Session()\n",
    "bucket_name = sess.default_bucket()\n",
    "default_bucket_prefix = sess.default_bucket_prefix\n",
    "default_bucket_prefix_path = \"\"\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    default_bucket_prefix_path = f\"/{default_bucket_prefix}\"\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "region = sess.boto_region_name\n",
    "\n",
    "iam_client = boto3.client(\"iam\")\n",
    "sts_client = boto3.client(\"sts\")\n",
    "sm_client = boto3.client(\"sagemaker\")\n",
    "account_id = sts_client.get_caller_identity()[\"Account\"]\n",
    "tracking_server_name = \"my-setup-test3\"\n",
    "mlflow_role_name = \"mlflow-test3\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c292837-353c-4c3c-91b9-3088e8d5a02b",
   "metadata": {},
   "source": [
    "## MLflow Permissions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6bae350-030f-4ecf-8380-5b11b73b5806",
   "metadata": {},
   "source": [
    "### IAM Role for the MLflow Tracking Server\n",
    "\n",
    "To run the next cell, make sure the IAM role used while running this notebook has permission to create an IAM Role. \n",
    "The `iam:CreateRole`, `iam:CreatePolicy`, `iam:ListPolicies`, and `iam:AttachRolePolicy` action must be allowed by the notebook execution role's policy.\n",
    "\n",
    "If you are running this notebook from SageMaker Studio, you can update your notebook execution role through the following steps: \n",
    "\n",
    "1. Navigate to the AWS Console and select the Domain you are using\n",
    "2. Under the Domain, select the User Profile you are using. You will see the Execution Role listed there.\n",
    "3. Navigate to the IAM Console, search for the Execution Role under \"Roles\", and update your role with a policy that allows the `iam:CreateRole`, `iam:CreatePolicy`, `iam:ListPolicies`, and `iam:AttachRolePolicy` actions. \n",
    "\n",
    "If you are not using a SageMaker Studio Notebook, confirm that the role you have used to configure your AWS CLI has appropriate permissions to create an IAM role and attach a policy to it. \n",
    "\n",
    "Here is an example of an inline policy you can add to your role - \n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"Statement1\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"iam:ListPolicies\",\n",
    "                \"iam:CreatePolicy\",\n",
    "                \"iam:CreateRole\",\n",
    "                \"iam:AttachRolePolicy\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"*\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c0ad98-f237-4bfd-b134-40b46ebfa81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow_trust_policy = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Principal\": {\"Service\": [\"sagemaker.amazonaws.com\"]},\n",
    "            \"Action\": \"sts:AssumeRole\",\n",
    "        }\n",
    "    ],\n",
    "}\n",
    "\n",
    "# Create role for MLflow\n",
    "mlflow_role = iam_client.create_role(\n",
    "    RoleName=mlflow_role_name, AssumeRolePolicyDocument=json.dumps(mlflow_trust_policy)\n",
    ")\n",
    "mlflow_role_arn = mlflow_role[\"Role\"][\"Arn\"]\n",
    "\n",
    "# Create policy for S3 and SageMaker Model Registry\n",
    "sm_s3_model_registry_policy = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"s3:Get*\",\n",
    "                \"s3:Put*\",\n",
    "                \"s3:List*\",\n",
    "                \"sagemaker:AddTags\",\n",
    "                \"sagemaker:CreateModelPackageGroup\",\n",
    "                \"sagemaker:CreateModelPackage\",\n",
    "                \"sagemaker:UpdateModelPackage\",\n",
    "                \"sagemaker:DescribeModelPackageGroup\",\n",
    "            ],\n",
    "            \"Resource\": \"*\",\n",
    "        }\n",
    "    ],\n",
    "}\n",
    "\n",
    "mlflow_s3_sm_model_registry_iam_policy = iam_client.create_policy(\n",
    "    PolicyName=\"mlflow-s3-sm-model-registry\", PolicyDocument=json.dumps(sm_s3_model_registry_policy)\n",
    ")\n",
    "mlflow_s3_sm_model_registry_iam_policy_arn = mlflow_s3_sm_model_registry_iam_policy[\"Policy\"][\"Arn\"]\n",
    "\n",
    "# Attach the policy to the MLflow role\n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=mlflow_role_name, PolicyArn=mlflow_s3_sm_model_registry_iam_policy_arn\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e2744c-c1b4-4444-9e8f-fbf1315a71a4",
   "metadata": {},
   "source": [
    "Note that your SageMaker execution role should have the following permissions to call Mlflow REST APIs:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",    \n",
    "    \"Statement\": [        \n",
    "        {            \n",
    "            \"Effect\": \"Allow\",            \n",
    "            \"Action\": [\n",
    "                \"sagemaker-mlflow:*\",\n",
    "                \"sagemaker:CreateMlflowTrackingServer\",\n",
    "                \"sagemaker:UpdateMlflowTrackingServer\",\n",
    "                \"sagemaker:DeleteMlflowTrackingServer\",\n",
    "                \"sagemaker:StartMlflowTrackingServer\",\n",
    "                \"sagemaker:StopMlflowTrackingServer\",\n",
    "                \"sagemaker:CreatePresignedMlflowTrackingServerUrl\"\n",
    "            ],            \n",
    "            \"Resource\": \"*\"        \n",
    "        }        \n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade88b9a-961a-4ced-9320-e56d7e9cf3eb",
   "metadata": {},
   "source": [
    "## Create MLflow Tracking Server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d496f9b-4493-4ab2-9d35-8d4ec0f79620",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.create_mlflow_tracking_server(\n",
    "    TrackingServerName=tracking_server_name,\n",
    "    ArtifactStoreUri=f\"s3://{bucket_name}{default_bucket_prefix_path}/{tracking_server_name}\",\n",
    "    TrackingServerSize=\"Small\",\n",
    "    MlflowVersion=\"2.13.2\",\n",
    "    RoleArn=mlflow_role_arn,\n",
    "    AutomaticModelRegistration=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d535f6-6dd2-4c5c-99e3-8b428c052c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracking_server_arn = (\n",
    "    f\"arn:aws:sagemaker:{region}:{account_id}:mlflow-tracking-server/{tracking_server_name}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddff09d7-73aa-4f77-b437-1e8c05c59ea2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.describe_mlflow_tracking_server(TrackingServerName=tracking_server_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c50a30-89e4-4ea9-8fe8-df15a2f7726e",
   "metadata": {},
   "source": [
    "Install the MLflow SDK and our MLflow AWS Plugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2232f516-f23c-4c0d-ada2-933a45fea6e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet mlflow==2.13.2 sagemaker-mlflow==0.1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073d12e9-b91e-4c0c-93d1-8cae66648e49",
   "metadata": {},
   "source": [
    "## MLflow tracking test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad90cde7-9de2-4df7-80a5-010165edafce",
   "metadata": {},
   "source": [
    "Connect to tracking server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7a43ce7-3e9a-4b47-b051-9f59522ee43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.set_tracking_uri(tracking_server_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9197fca-6370-4f91-a52f-440ef5b22484",
   "metadata": {},
   "source": [
    "Log a metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab5d5df-c1a8-4a2b-89e1-52d36d630f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with mlflow.start_run():\n",
    "    mlflow.log_metric(\"foo\", 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d603ef2f-9c42-4ef2-896e-73ab1eaa6ace",
   "metadata": {},
   "source": [
    "See results in MLflow UI. You can either launch the MLflow UI from within SageMaker Studio, or generate a pre-signed URL like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0690708f-976c-472e-8e4d-281aa163e9aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.create_presigned_mlflow_tracking_server_url(TrackingServerName=tracking_server_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f739f1a-2a97-4cc5-bb6b-bc59e4111d0f",
   "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",
    "\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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.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/ml_ops|sm-mlflow_setup|sm-mlflow_setup.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}