update setup instructions

Author: BassemHalim

v3-examples/ml-ops-examples/v3-feature-store-lake-formation.ipynb

@@ -7,6 +7,25 @@
     "# Feature Group Lake Formation Governance "
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This Demo builds on sagemaker v3 features which is not out yet so we need to tell pip to use local version \n",
+    "\n",
+    "1. install `pyenv` following the instructions here https://github.com/pyenv/pyenv?tab=readme-ov-file#a-getting-pyenv\n",
+    "2. setup your shell env for pyenv https://github.com/pyenv/pyenv?tab=readme-ov-file#b-set-up-your-shell-environment-for-pyenv\n",
+    "3. install and activate python for example (this code was tested on python 3.10.14)\n",
+    "```\n",
+    "pyenv install 3.10.14\n",
+    "pyenv virtualenv 3.10.14 py3.10.14\n",
+    "pyenv activate py3.10.14\n",
+    "```\n",
+    "3. You will then need to run `pip install -e ./sagemaker-mlops` before starting the jupyter notebook to use the local packages\n",
+    "4. `pip install jupyter notebook`\n",
+    "5. `jupyter notebook`"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -230,21 +249,43 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "This Demo builds on sagemaker v3 features which is not out yet so we need to tell pip to use local version \n",
+    "import boto3\n",
+    "import os\n",
     "\n",
-    "1. install `pyenv` https://github.com/pyenv/pyenv?tab=readme-ov-file#a-getting-pyenv\n",
-    "2. install and activate python for example (this code was tested on python 3.10.14)\n",
-    "```\n",
-    "pyenv install 3.10.14\n",
-    "pyenv virtualenv 3.10.14 py3.10.14\n",
-    "pyenv activate py3.10.14\n",
-    "```\n",
-    "3. You will then need to run `pip install -e .` before starting the jupyter notebook to use the local packages\n",
-    "4. `pip install jupyter notebook`\n",
-    "5. `jupyter notebook`"
+    "boto3.DEFAULT_SESSION = None\n",
+    "def assume_role(role_arn, session_name=\"AssumedRoleSession\"):\n",
+    "    \"\"\"\n",
+    "    Assume an AWS IAM role and return temporary credentials.\n",
+    "    \n",
+    "    Args:\n",
+    "        role_arn: The ARN of the role to assume\n",
+    "        session_name: A name for the assumed role session\n",
+    "    \n",
+    "    Returns:\n",
+    "        A boto3 session with the assumed role credentials\n",
+    "    \"\"\"\n",
+    "    sts_client = boto3.client('sts')\n",
+    "    \n",
+    "    response = sts_client.assume_role(\n",
+    "        RoleArn=role_arn,\n",
+    "        RoleSessionName=session_name\n",
+    "    )\n",
+    "    \n",
+    "    credentials = response['Credentials']\n",
+    "    \n",
+    "    # Create a new session with the temporary credentials\n",
+    "    session = boto3.Session(\n",
+    "        aws_access_key_id=credentials['AccessKeyId'],\n",
+    "        aws_secret_access_key=credentials['SecretAccessKey'],\n",
+    "        aws_session_token=credentials['SessionToken']\n",
+    "    )\n",
+    "    \n",
+    "    return session"
    ]
   },
   {
@@ -287,22 +328,26 @@
    "source": [
     "# Use SageMaker session to get default bucket and execution role\n",
     "\n",
-    "boto_session = boto3.Session()\n",
-    "new_sagemaker_session = SageMakerSession(boto_session=boto_session)\n",
-    "sagemaker_session = SageMakerSession()\n",
-    "S3_BUCKET = sagemaker_session.default_bucket()\n",
-    "REGION = sagemaker_session.boto_session.region_name\n",
-    "\n",
     "# Execution role (for running this notebook)\n",
-    "EXECUTION_ROLE_ARN = get_execution_role(sagemaker_session)\n",
+    "EXECUTION_ROLE_ARN = 'arn:aws:iam::<account id>:role/<role>'\n",
+    "\n",
     "\n",
     "# Offline store role (dedicated role for Feature Store S3 access)\n",
     "# Replace with your dedicated offline store role ARN\n",
     "# https://docs.aws.amazon.com/sagemaker/latest/dg/feature-store-adding-policies.html\n",
-    "OFFLINE_STORE_ROLE_ARN = \"arn:aws:iam::<aws account id>:role/<role name>\"\n",
+    "OFFLINE_STORE_ROLE_ARN = \"arn:aws:iam::<account id>:role/<role>\"\n",
+    "\n",
+    "\n",
+    "boto_session = assume_role(EXECUTION_ROLE_ARN)\n",
+    "sagemaker_session = SageMakerSession(boto_session=boto_session)\n",
+    "sts = boto_session.client('sts')\n",
+    "\n",
+    "S3_BUCKET = sagemaker_session.default_bucket()\n",
+    "REGION = sagemaker_session.boto_session.region_name\n",
+    "\n",
     "\n",
     "print(f\"S3 Bucket: {S3_BUCKET}\")\n",
-    "print(f\"Execution Role ARN: {EXECUTION_ROLE_ARN}\")\n",
+    "print(f\"Execution Role ARN:{sts.get_caller_identity()['Arn']}\")\n",
     "print(f\"Offline Store Role ARN: {OFFLINE_STORE_ROLE_ARN}\")\n",
     "print(f\"Region: {REGION}\")"
    ]
@@ -572,12 +617,10 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "fg_workflow2 = FeatureGroup.refresh()"
+    "Here you can optionally assume a different role with just datalake permissions and describe Feature Group permission"
    ]
   },
   {
@@ -593,10 +636,9 @@
     "print(\"  2. Grant permissions to execution role\")\n",
     "print(\"  3. Revoke IAMAllowedPrincipal permissions\")\n",
     "print()\n",
-    "\n",
+    "fg_workflow2 = FeatureGroup.get(FG_NAME_WORKFLOW2)\n",
     "result = fg_workflow2.enable_lake_formation( # new method\n",
     "    use_service_linked_role=True,\n",
-    "    session=boto_session\n",
     ")\n",
     "\n",
     "print(f\"\\nLake Formation setup results:\")\n",
@@ -645,7 +687,7 @@
     "print(f\"Status: {fg_workflow2.feature_group_status}\")\n",
     "print(f\"ARN: {fg_workflow2.feature_group_arn}\")\n",
     "DataCatalogTable = fg_workflow2.offline_store_config.data_catalog_config.table_name\n",
-    "print(f\"{DataCatalogTable}\")"
+    "print(f\"Table Name: {DataCatalogTable}\")"
    ]
   },
   {
@@ -715,7 +757,7 @@
    "outputs": [],
    "source": [
     "# Uncomment to delete the Feature Groups\n",
-    "# cleanup_feature_group(fg_workflow1)\n",
+    "cleanup_feature_group(fg_workflow1)\n",
     "cleanup_feature_group(fg_workflow2)"
    ]
   },