{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CloudFormation script to automate cross-account resource registration on central Model Registry\n",
    "\n",
    "* Please reach out to Yuyao Zhang ozhang@amazon.com or Melanie Li mmelli@amazon.com for any issue or questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install boto3 sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The role for CloudFormation requires `SageMakerFullAccess` and following additional permissions:\n",
    "```\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"VisualEditor0\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"iam:GetRole\",\n",
    "                \"events:DescribeRule\",\n",
    "                \"events:PutRule\",\n",
    "                \"iam:DeletePolicy\",\n",
    "                \"iam:CreateRole\",\n",
    "                \"iam:DeleteRole\",\n",
    "                \"iam:AttachRolePolicy\",\n",
    "                \"iam:PutRolePolicy\",\n",
    "                \"iam:CreatePolicy\",\n",
    "                \"events:PutTargets\",\n",
    "                \"events:DeleteRule\",\n",
    "                \"iam:PassRole\",\n",
    "                \"iam:DetachRolePolicy\",\n",
    "                \"iam:DeleteRolePolicy\",\n",
    "                \"events:RemoveTargets\",\n",
    "                \"iam:GetRolePolicy\",\n",
    "                \"iam:DeletePolicyVersion\"\n",
    "            ],\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting prod-account-template.yaml\n"
     ]
    }
   ],
   "source": [
    "%%writefile prod-account-template.yaml\n",
    "---\n",
    "AWSTemplateFormatVersion: '2010-09-09'\n",
    "\n",
    "Description: >-\n",
    "  Trigger SageMaker Lineage Creation in Root account.\n",
    "Parameters:\n",
    "  RootAccountNumber:\n",
    "    Type: String\n",
    "    Description: Root acount where events should be sent to.\n",
    "        \n",
    "Resources:\n",
    "\n",
    "  #### SECTION: Triggering the Lambda Automatically\n",
    "  # Until SageMaker directly supports EventBridge events on 'apps' (as notebook instance statuses already are), we can\n",
    "  # work around this by triggering the function via a CloudTrail. \n",
    "\n",
    "  CloudWatchEventRole:\n",
    "    Type: AWS::IAM::Role\n",
    "    Properties:\n",
    "      AssumeRolePolicyDocument:\n",
    "        Version: '2012-10-17'\n",
    "        Statement:\n",
    "          - Effect: Allow\n",
    "            Principal:\n",
    "              Service:\n",
    "                - events.amazonaws.com\n",
    "            Action: sts:AssumeRole\n",
    "      Path: /\n",
    "      Policies:\n",
    "        - PolicyName: AllowPutEvents\n",
    "          PolicyDocument:\n",
    "            Version: '2012-10-17'\n",
    "            Statement:\n",
    "              - Effect: Allow\n",
    "                Action: 'events:PutEvents'\n",
    "                Resource: !Sub 'arn:aws:events:${AWS::Region}:${RootAccountNumber}:event-bus/default'\n",
    "\n",
    "  # Finally, the rule controls how CloudTrail-logged events trigger the Lambda\n",
    "  CloudWatchEventRule:\n",
    "    Type: 'AWS::Events::Rule'\n",
    "    Properties:\n",
    "      EventPattern:\n",
    "        source:\n",
    "          - aws.sagemaker\n",
    "        detail-type:\n",
    "          - 'AWS API Call via CloudTrail'\n",
    "        detail:\n",
    "          eventSource:\n",
    "            - sagemaker.amazonaws.com\n",
    "          eventName:\n",
    "            - CreateModel\n",
    "          requestParameters:\n",
    "            containers:\n",
    "              modelPackageName:\n",
    "              - exists: true\n",
    "      Targets:\n",
    "        - Arn: !Sub 'arn:aws:events:${AWS::Region}:${RootAccountNumber}:event-bus/default'\n",
    "          Id: api-based-createmodel-lineage\n",
    "          RoleArn: !GetAtt CloudWatchEventRole.Arn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_account_id = 607162686141"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"StackId\": \"arn:aws:cloudformation:ap-southeast-2:043207074741:stack/SageMakerLineageExtention/5af09230-dc2c-11ec-9ded-02a200717f74\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "!aws cloudformation create-stack \\\n",
    "  --stack-name SageMakerLineageExtention\\\n",
    "  --template-body file://./prod-account-template.yaml \\\n",
    "  --parameters ParameterKey=RootAccountNumber,ParameterValue=$root_account_id \\\n",
    "  --role-arn arn:aws:iam::043207074741:role/cfntest \\\n",
    "  --capabilities CAPABILITY_IAM  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sess = sagemaker.session.Session()\n",
    "\n",
    "modelPkg = sagemaker.model.ModelPackage(name='cross-account-test-00',\n",
    "                     model_package_arn=\"arn:aws:sagemaker:ap-southeast-2:607162686141:model-package/woodside-new-p-6dvnevzoh2mt/2\", \n",
    "                     role=\"arn:aws:iam::043207074741:role/service-role/AmazonSageMaker-ExecutionRole-20220522T213009\",\n",
    "                     sagemaker_session=sess)\n",
    "\n",
    "# Create SageMaker Model now\n",
    "modelPkg._create_sagemaker_model()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}