AWS Lex Conversational FAQ Demo: CDK code to create lambdas, lex bot and sagemaker endpoint. User uploads file to create index

Author: mhenkelw-aws-lex

DEVELOPMENT.md

@@ -0,0 +1,45 @@
+
+### Project Structure within src/lex-gen-ai-demo-cdk:
+```
+AWSLexKoiosBlogDemo/src/lex-gen-ai-demo-cdk
+- app.py
+- cdk.json
+- endpoint_handler.py
+- upload_file_to_s3.py
+- shutdown_endpoint.py
+- index-creation-docker-image/
+   - index_creation_app.py
+   - Dockerfile
+   - index_creation_requirements.txt
+- lex_gen_ai_demo_cdk_files/
+   - __init__.py
+   - lex_gen_ai_demo_file_stack.py
+- lex-gen-ai-demo-docker-image/
+   - runtime_lambda_app.py
+   - Dockerfile
+   - runtime_lambda_requirements.txt
+- requirements.txt
+- source.bat
+```
+
+## Common Errors & Troubleshooting
+
+### "ValueError: Must setup local AWS configuration with a region supported by SageMaker."
+Solution: You must set an aws region with `export AWS_DEFAULT_REGION=<your-region>`
+
+### Error creating role
+```
+botocore.exceptions.ClientError: An error occurred (AccessDenied) when calling the CreateRole operation: User: <user-arn> is not authorized to perform: iam:CreateRole on resource: <role-arn> because no identity-based policy allows the iam:CreateRole action
+```
+Solution: you must ensure the Iam role you are using has sufficient permissions to create Iam roles
+
+### Error LexGenAIDemoFilesStack:  fail: docker push <IMAGE> exited with error code 1: tag does not exist
+Issue: Error while building the image. Here are some common ones 
+
+#### Error processing tar file(exit status 1): write /path/libcublas.so.11: no space left on device
+Issue: Docker has run out of memory due to too many images
+Solution: Delete unused images in the Docker application and then [prune docker](https://docs.docker.com/config/pruning/) in command line 
+
+#### ConnectionResetError: [Errno 104] Connection reset by peer
+Issue: Pip issue
+Solution: Clear pip cache (`python3 -m pip cache purge`) and run again

README.md

@@ -1,17 +1,124 @@
-## My Project
+# AWS Lex Conversational FAQ Demo
 
-TODO: Fill this README out!
+Demonstration of LLM integration into a lex bot using Lambda codehooks and a Sagemaker endpoint.
 
-Be sure to:
+![Diagram](diagram.png)
 
-* Change the title in this README
-* Edit your repository description on GitHub
+### What resources will be created?
+This CDK code will create the following:
+   - 1 Sagemaker endpoint hosting a model (falcon-7b-instruct on ml.g5.8xlarge by default but this is configurable)
+   - 1 Lex bot
+   - 2 S3 buckets (one for your uploaded source, one for the created index)
+   - 2 Lambda functions (one to ingest the source and create an image, one to be invoked as codehook during lambda and provide an FAQ answer when needed)
+   - 1 Event listener attached to an S3 bucket to call the index creation lambda automatically when a file is uploaded
+   - 2 Iam roles (one for the lex bot to call lambda, one for the lambdas to call sagemaker and S3)
 
-## Security
+## Requirements
 
-See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information.
+### AWS setup
+**Region**
 
-## License
+If you have not yet run `aws configure` and set a default region, you must do so, or you can also run `export AWS_DEFAULT_REGION=<your-region>` 
 
-This library is licensed under the MIT-0 License. See the LICENSE file.
+**Authorization**
 
+You must use a role that has sufficient permissions to create Iam roles, as well as cloudformation resources
+
+### Python >3.7
+Make sure you have [python3](https://www.python.org/downloads/) installed at a version >=3.7.x
+
+### Docker
+Make sure you have [Docker](https://www.docker.com/products/docker-desktop/) installed on your machine and running in the background 
+
+### AWS CDK
+Make sure you have the [AWS CDK](https://docs.aws.amazon.com/cdk/v2/guide/getting_started.html#getting_started_install) installed on your machine
+
+
+## Setup
+
+### Set up virtual enviroment and gather packages
+
+```
+cd src/lex-gen-ai-demo-cdk-files
+```
+
+Install the required dependencies (aws-cdk-lib and constructs) into your Python environment 
+```
+pip install -r requirements.txt
+```
+
+### Gather and deploy resources with the CDK
+
+First synthesize, which executes the application, defines which resources will be created and translates this into a cloudformation template
+```
+cdk synth
+```
+Now bootstrap, which provisions the resources you'll use when deploying the application
+```
+cdk bootstrap
+```
+and deploy with
+```
+cdk deploy LexGenAIDemoFilesStack
+```
+
+The deployment will create a lex bot and S3 buckets and will dockerize the code in the koios_cdk_files/koios-docker-image directory and push that image to ECR so it can run in Lambda. Dont worry if this step takes a long time while pushing to ECR, we are bundling up two docker images and uploading them so it will take some time.
+
+## Usage
+Once all the resources are created after `cdk deploy` finishes running you must upload a .PDF or .txt file at least once so an index can be created. You can use our upload script `upload_file_to_s3.py path/to/your/file` or you can navigate to the S3 console and manually upload a file. On upload the ingestion lambda will read the file and create an embedding which it will upload to the other S3 bucket. Now that an embedding exists you can go to your bot and begin using it. If you want to update the embedding you can upload a new file and a new embedding will overwrite the old embedding. Once you have a new embedding you must restart the runtime lambda function for it to start using the new embedding. 
+
+Note, the first time the embedding lambda and the runtime lambda are called the latency will be much slower as it must load resources and save them in the lambda enviroment. Once loaded these resources will stay in the enviroment as long as the ECR image is not deleted. This means your first request will be slow but after that it will speed up now that the resources are cached.
+
+### Uploading files 
+Now, you have to upload your source file so the indexing lambda can create an index for the runtime lambda function to use. You can use our script with any PDF or .txt file by running
+```
+python3 upload_file_to_s3.py path/to/your/file
+```
+or you can open the S3 bucket in the console and manually upload a file. On upload an index will automatically be generated.
+Note: If you upload a large file, the index will be large and the S3 read time on cold start may become large.
+
+Once you've uploaded your file, wait a little for your index to be created and then you can go into the Lex console and test your bot (no need to build your bot unless you've made changes after creation). The first time you create an index and the first time you query the bot it will take a little longer (around 90 seconds) as we need to load models and cache them in the lambda-ECR enviroment but after they are cached there is no need to download them and latency will be much faster. These resources will remain cached as long as the ECR image is not deleted. Additionally for better cold start performance you can an instance for your runtime lambda function. There are directions to do so below.
+
+### Configurations
+
+🚨 **Remember to shut down your endpoint if you're done using it!** 🚨
+
+We have provided a script to deactivate an endpoint and endpoint configuration with whatever name is in the endpoint creation script. To run:
+```
+python3 shut_down_endpoint.py
+```
+
+#### Custom model and instance type configuration:
+
+The function `create_endpoint_from_HF_image()` is called in `app.py`. This function acceptst the following arguments:
+ - hf_model_id (required): For the purposes of the demo we have this set to [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b). You can find any https://huggingface.co/ and feed it in
+ - instance_type (optional, default is ml.g5.8xlarge): If you dont give an argument we'll use ml.g5.8xlarge. You can use any endpoint [sage instance type](https://aws.amazon.com/sagemaker/pricing/)
+ - endpoint_name (optional, default is whatever SAGEMAKER_ENDPOINT_NAME is set to in the file endpoint_handler.py): You can give your endpoint a custom name. It is recomended that you don't do this but if you do, you have to change it in the lamdba images (constant is called ENDPOINT_NAME in index_creation_app.py and runtime_lambda_app.py)
+ - number_of_gpu (optional, default is 1): Set this to any number of GPUs the hardware you chose allows. 
+
+ If you have in invalid configuration the endpoint will fail to create. You can see the specific error in the cloudwatch logs. If you fail creation you can run `python3 shut_down_endpoint.py` to clean up the endpoint but if you do so manually in the console **you must delete both the endpoint and the endpoint configuration**
+
+#### Fruther configuration
+If you would like to further configure the endpoint you can change the specific code in `endpoint_handler.py`
+
+The LLM is hosted on a sagemaker endpoint and deployed as a sagemaker [HuggingFaceModel](https://sagemaker.readthedocs.io/en/stable/frameworks/huggingface/sagemaker.huggingface.html). We are also using a huggingface model image. You can read more about it [here](https://aws.amazon.com/blogs/machine-learning/announcing-the-launch-of-new-hugging-face-llm-inference-containers-on-amazon-sagemaker/). For further model configuration you can read about sagemaker model deployments [here](https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints-deployment.html).
+
+For our indexing and retrieval we are using [llama-index](https://github.com/jerryjliu/llama_index). If you would like to configure the index retriever you can do so in the `runtime_lambda_app.py` file in the `VectorIndexRetriever` object on line 70. If you want to update index creation you can update the constants defined at the top of the index creation and runtime lambdas (`index_creation_app.py` and `runtime_lambda_app.py`). Make sure to familiarize yourself with [llama-index terms](https://gpt-index.readthedocs.io/en/latest/guides/tutorials/terms_definitions_tutorial.html) and the [llama-index prompthelper](https://gpt-index.readthedocs.io/en/latest/reference/service_context/prompt_helper.html) for best results.
+
+### Tips for best results
+
+**Keep your lambda perpetually warm by provisioning an instance for the runtime lambda (lex-codehook-fn)**
+
+Go to Lambda console > select the function lex-codehook-fn
+
+Versions > Publish new version
+
+Under this version 
+   - Provisioned Concurrency > set value to 1
+   - Permissions > Resource based policy statements > Add Permissions > AWS Service > Other, your-policy-name, lexv2.amazonaws.com, your-lambda-arn, lamdba:InvokeFunction
+
+Go to your Lex Bot (LexGenAIDemoBotCfn)
+
+Aliases > your-alias > your-language > change lambda function version or alias > change to your-version
+
+This will keep an instance running at all times and keep your lambda ready so that you wont have cold start latency. This will cost a bit extra (https://aws.amazon.com/lambda/pricing/) so use thoughfully. 

diagram.png

@@ -0,0 +1,14 @@
+
+import aws_cdk as cdk
+
+from lex_gen_ai_demo_cdk_files.lex_gen_ai_demo_cdk_files_stack import LexGenAIDemoFilesStack
+from endpoint_handler import create_endpoint_from_HF_image
+
+# create_endpoint_from_HF_image(hf_model_id, instance_type="ml.g5.8xlarge", endpoint_name=SAGEMAKER_ENDPOINT_NAME, number_of_gpu=1)
+# You can run with no arguments to get default values of google/flan-t5-xxl on ml.g5.8xlarge, or pass in your own arguments
+create_endpoint_from_HF_image(hf_model_id="tiiuae/falcon-7b-instruct")
+
+app = cdk.App()
+filestack = LexGenAIDemoFilesStack(app, "LexGenAIDemoFilesStack")
+
+app.synth()

src/lex-gen-ai-demo-cdk/app.py

@@ -0,0 +1,51 @@
+{
+    "app": "python3 app.py",
+    "watch": {
+      "include": [
+        "**"
+      ],
+      "exclude": [
+        "README.md",
+        "cdk*.json",
+        "requirements*.txt",
+        "source.bat",
+        "**/__init__.py",
+        "python/__pycache__",
+        "tests"
+      ]
+    },
+    "context": {
+      "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
+      "@aws-cdk/core:checkSecretUsage": true,
+      "@aws-cdk/core:target-partitions": [
+        "aws",
+        "aws-cn"
+      ],
+      "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true,
+      "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true,
+      "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true,
+      "@aws-cdk/aws-iam:minimizePolicies": true,
+      "@aws-cdk/core:validateSnapshotRemovalPolicy": true,
+      "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true,
+      "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true,
+      "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true,
+      "@aws-cdk/aws-apigateway:disableCloudWatchRole": true,
+      "@aws-cdk/core:enablePartitionLiterals": true,
+      "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true,
+      "@aws-cdk/aws-iam:standardizedServicePrincipals": true,
+      "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true,
+      "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true,
+      "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true,
+      "@aws-cdk/aws-route53-patters:useCertificate": true,
+      "@aws-cdk/customresources:installLatestAwsSdkDefault": false,
+      "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true,
+      "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true,
+      "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true,
+      "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true,
+      "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true,
+      "@aws-cdk/aws-redshift:columnId": true,
+      "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true,
+      "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true,
+      "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true
+    }
+  }

src/lex-gen-ai-demo-cdk/cdk.json

@@ -0,0 +1,99 @@
+import json
+import boto3
+import time
+from sagemaker.huggingface import get_huggingface_llm_image_uri
+from sagemaker.huggingface import HuggingFaceModel
+
+# get image from huggingface
+llm_image = get_huggingface_llm_image_uri(
+  "huggingface",
+  version="0.8.2"
+)
+
+assume_role_policy_document = json.dumps({
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+        "Effect": "Allow",
+        "Principal": {
+            "Service": [
+                "sagemaker.amazonaws.com",
+                "ecs.amazonaws.com"
+            ]
+        },
+        "Action": "sts:AssumeRole"
+        }
+    ]
+})
+
+# editable to whatever you want your endpoint and role to be. You can use an existing role or a new one
+# IMPORTANT: make sure your lambda endpoint name in lambda_app.py is consisitent if you change it here
+SAGEMAKER_IAM_ROLE_NAME = 'Sagemaker-Endpoint-Creation-Role'
+SAGEMAKER_ENDPOINT_NAME = "huggingface-pytorch-sagemaker-endpoint"
+
+# Create role and give sagemaker permissions
+def get_iam_role(role_name=SAGEMAKER_IAM_ROLE_NAME):
+    iam_client = boto3.client('iam')
+
+    try: 
+        role = iam_client.get_role(RoleName=role_name)
+        role_arn = role['Role']['Arn']
+        print(f"Role {role_arn} found!")
+        return role_arn
+    
+    except:
+        role_arn = iam_client.create_role(
+            RoleName=SAGEMAKER_IAM_ROLE_NAME,
+            AssumeRolePolicyDocument=assume_role_policy_document
+            )['Role']['Arn']
+
+        time.sleep(10) # give the policy some time to properly create
+
+        response = iam_client.attach_role_policy(
+            PolicyArn='arn:aws:iam::aws:policy/AmazonSageMakerFullAccess',
+            RoleName=SAGEMAKER_IAM_ROLE_NAME,
+        )
+        print(f"Creating {role_arn}")
+        time.sleep(20) # give iam time to let the role create
+        return role_arn
+
+
+# Define Model and Endpoint configuration parameter
+
+health_check_timeout = 300
+trust_remote_code = True
+
+# Create sagemaker endpoint, default values are flan t5 xxl in a g5.8xl instance
+def create_endpoint_from_HF_image(hf_model_id, instance_type="ml.g5.8xlarge", endpoint_name=SAGEMAKER_ENDPOINT_NAME, number_of_gpu=1):
+    sagemaker_client = boto3.client('sagemaker')
+
+    try: # check if endpoint already existst
+        sagemaker_client.describe_endpoint(EndpointName=SAGEMAKER_ENDPOINT_NAME)
+        print(f"Endpoint with name {SAGEMAKER_ENDPOINT_NAME} found!")
+        return
+    
+    except:
+        print(f"Creating endpoint with model{hf_model_id} on {instance_type}...")
+
+        # create HuggingFaceModel with the image uri
+        llm_model = HuggingFaceModel(
+            role=get_iam_role(),
+            image_uri=llm_image,
+            env={
+                'HF_MODEL_ID': hf_model_id,
+                'SM_NUM_GPUS': json.dumps(number_of_gpu),
+                'HF_MODEL_TRUST_REMOTE_CODE': json.dumps(trust_remote_code)
+            }
+        )
+
+        # Deploy model to an endpoint
+        # https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy
+        llm = llm_model.deploy(
+            endpoint_name=endpoint_name,
+            initial_instance_count=1,
+            instance_type=instance_type,
+            # volume_size=400, # If using an instance with local SSD storage, volume_size must be None, e.g. p4 but not p3
+            container_startup_health_check_timeout=health_check_timeout  # 10 minutes to be able to load the model
+        )
+
+        print(f"\nEndpoint created ({endpoint_name})")

src/lex-gen-ai-demo-cdk/endpoint_handler.py

@@ -0,0 +1,14 @@
+FROM public.ecr.aws/lambda/python:3.8
+
+COPY index_creation_requirements.txt  .
+RUN  pip3 install -r index_creation_requirements.txt --target "${LAMBDA_TASK_ROOT}"
+ 
+# Copy function code
+COPY *.py ${LAMBDA_TASK_ROOT}
+ 
+# Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile)
+CMD [ "index_creation_app.handler" ]
+
+# Set cache to a location lambda can write to
+ENV TRANSFORMERS_CACHE="/tmp/TRANSFORMERS_CACHE"
+