set up the finetune jobs endpoints

botisan-ai · Sep 2, 2021 · c428339 · c428339
1 parent 8ba8f1a
commit c428339
Show file tree

Hide file tree

Showing 10 changed files with 318 additions and 1 deletion.
diff --git a/functions/start_finetune_job.py b/functions/start_finetune_job.py
@@ -0,0 +1,68 @@
+import os
+import json
+import traceback
+import openai
+import boto3
+from typing import Any, Dict
+from dotenv import dotenv_values
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+config = {
+    **dotenv_values(os.path.join(os.path.dirname(__file__), '..', '.env')),
+    **os.environ,
+}
+
+openai.api_key = config.get('OPENAI_API_KEY')
+BUCKET_NAME = config.get('BUCKET_NAME')
+
+
+def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
+    try:
+        uid = event.get('requestContext', {}).get(
+            'authorizer', {}).get('principalId')
+
+        body = event.get('body')
+
+        if body is None:
+            return {
+                'statusCode': 400,
+                'body': json.dumps({'message': 'please provide POST body'}),
+                'headers': {
+                    'Content-Type': 'application/json',
+                },
+            }
+
+        body = json.loads(body)
+        key = body.get('key')
+
+        s3_resource = boto3.resource('s3')
+
+        s3_object = s3_resource.Object(BUCKET_NAME, key)
+        s3_response = s3_object.get()
+        file_response = openai.File.create(file=s3_response['Body'], purpose='fine-tune')
+        file_id = file_response.get('id', '')
+        finetune_response = openai.FineTune.create(training_file=file_id)
+        new_s3_object = s3_resource.Object(BUCKET_NAME, f'{uid}/{file_id}')
+        new_s3_object.copy_from(CopySource=key)
+        s3_object.delete()
+
+        return {
+            'statusCode': 200,
+            'body': json.dumps(finetune_response),
+            'headers': {
+                'Content-Type': 'application/json',
+            },
+        }
+    except Exception:
+        traceback.print_exc()
+        return {
+            'statusCode': 400,
+            'body': json.dumps({'message': 'error occurred, please check logs'}),
+            'headers': {
+                'Content-Type': 'application/json',
+            },
+        }
+
+
+if __name__ == '__main__':
+    print(config)
diff --git a/functions/upload_finetune_file.py b/functions/upload_finetune_file.py
@@ -0,0 +1,48 @@
+import os
+import json
+import uuid
+import traceback
+import boto3
+from typing import Any, Dict
+from dotenv import dotenv_values
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+config = {
+    **dotenv_values(os.path.join(os.path.dirname(__file__), '..', '.env')),
+    **os.environ,
+}
+
+BUCKET_NAME = config.get('BUCKET_NAME')
+
+
+def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
+    try:
+        uid = event.get('requestContext', {}).get(
+            'authorizer', {}).get('principalId')
+
+        s3_client = boto3.client('s3')
+
+        object_key = f'{uid}/{str(uuid.uuid4())}.jsonl'
+
+        response = s3_client.generate_presigned_post(BUCKET_NAME, object_key)
+
+        return {
+            'statusCode': 200,
+            'body': json.dumps(response),
+            'headers': {
+                'Content-Type': 'application/json',
+            },
+        }
+    except Exception:
+        traceback.print_exc()
+        return {
+            'statusCode': 400,
+            'body': json.dumps({'message': 'error occurred, please check logs'}),
+            'headers': {
+                'Content-Type': 'application/json',
+            },
+        }
+
+
+if __name__ == '__main__':
+    print(config)
diff --git a/package.json b/package.json
@@ -1,16 +1,21 @@
 {
   "name": "openai-serverless",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "OpenAI Serverless Functions",
   "main": "index.js",
   "author": "Xanthous Tech <hi@x-tech.io>",
   "license": "ISC",
   "private": true,
+  "scripts": {
+    "postinstall": "patch-package"
+  },
   "devDependencies": {
     "dotenv-cli": "^4.0.0",
+    "patch-package": "^6.4.7",
     "serverless": "^2.44.0",
     "serverless-domain-manager": "^5.1.0",
     "serverless-offline": "^7.0.0",
+    "serverless-plugin-additional-stacks": "^1.5.0",
     "serverless-python-requirements": "^5.1.1"
   }
 }
diff --git a/patches/serverless-plugin-additional-stacks+1.5.0.patch b/patches/serverless-plugin-additional-stacks+1.5.0.patch
@@ -0,0 +1,28 @@
+diff --git a/node_modules/serverless-plugin-additional-stacks/index.js b/node_modules/serverless-plugin-additional-stacks/index.js
+index d5f1379..0b1eff9 100644
+--- a/node_modules/serverless-plugin-additional-stacks/index.js
++++ b/node_modules/serverless-plugin-additional-stacks/index.js
+@@ -58,6 +58,7 @@ class AdditionalStacksPlugin {
+                 usage: 'Additional stack name to deploy',
+                 shortcut: 'k',
+                 required: false,
++                type: 'string',
+               },
+             },
+           },
+@@ -66,6 +67,7 @@ class AdditionalStacksPlugin {
+           'skip-additionalstacks': {
+             usage: 'Skip deploying additional stacks',
+             required: false,
++            type: 'boolean',
+           },
+         },
+       },
+@@ -81,6 +83,7 @@ class AdditionalStacksPlugin {
+                 usage: 'Additional stack name to remove',
+                 shortcut: 'k',
+                 required: false,
++                type: 'string',
+               },
+             },
+           },
diff --git a/resources/dynamodb.yml b/resources/dynamodb.yml
@@ -0,0 +1,26 @@
+stack:
+  StackName: ${self:custom.prefix}-dynamodb
+  Resources:
+    UsageTable:
+      Type: AWS::DynamoDB::Table
+      Properties:
+        TableName: ${self:custom.prefix}-usage
+        AttributeDefinitions:
+          - AttributeName: uid
+            AttributeType: S
+          - AttributeName: timestamp
+            AttributeType: N
+        KeySchema:
+          - AttributeName: uid
+            KeyType: HASH
+          - AttributeName: timestamp
+            KeyType: RANGE
+        ProvisionedThroughput:
+          ReadCapacityUnits: 5
+          WriteCapacityUnits: 5
+  Outputs:
+    UsageTableName:
+      Value:
+        Ref: UsageTable
+      Export:
+        Name: ${self:custom.prefix}-usage-table-name
diff --git a/resources/iam.yml b/resources/iam.yml
@@ -0,0 +1,21 @@
+# TODO: make these permissions more granular, this is too broad
+Resources:
+  lambdaRole:
+    Type: AWS::IAM::Role
+    Properties:
+      Description: Role for running proper actions in lambda functions.
+      RoleName: ${self:custom.prefix}-lambda-role
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: 'Allow'
+            Principal:
+              Service:
+                - 'lambda.amazonaws.com'
+            Action:
+              - 'sts:AssumeRole'
+      ManagedPolicyArns:
+        - 'arn:aws:iam::aws:policy/AmazonDynamoDBFullAccess'
+        - 'arn:aws:iam::aws:policy/AmazonSNSFullAccess'
+        - 'arn:aws:iam::aws:policy/CloudWatchLogsFullAccess'
+        - 'arn:aws:iam::aws:policy/AmazonS3FullAccess'
diff --git a/resources/openai_functions.yml b/resources/openai_functions.yml
@@ -46,3 +46,39 @@ extract_todo_list:
           type: token
   environment:
     OPENAI_API_KEY: ${env:OPENAI_API_KEY}
+start_finetune_job:
+  runtime: python3.8
+  handler: functions.start_finetune_job.handler
+  events:
+    - http:
+        path: finetune_jobs/start
+        method: post
+        cors: true
+        authorizer:
+          name: authorizer_function
+          resultTtlInSeconds: 0
+          identitySource: method.request.header.Authorization
+          identityValidationExpression: jwt .+
+          type: token
+  environment:
+    OPENAI_API_KEY: ${env:OPENAI_API_KEY}
+    BUCKET_NAME:
+      Fn::ImportValue: ${self:custom.prefix}-finetune-data-bucket-name
+upload_finetune_file:
+  runtime: python3.8
+  handler: functions.upload_finetune_file.handler
+  events:
+    - http:
+        path: finetune_jobs/upload
+        method: get
+        cors: true
+        authorizer:
+          name: authorizer_function
+          resultTtlInSeconds: 0
+          identitySource: method.request.header.Authorization
+          identityValidationExpression: jwt .+
+          type: token
+  environment:
+    OPENAI_API_KEY: ${env:OPENAI_API_KEY}
+    BUCKET_NAME:
+      Fn::ImportValue: ${self:custom.prefix}-finetune-data-bucket-name
diff --git a/resources/s3.yml b/resources/s3.yml
@@ -0,0 +1,13 @@
+stack:
+  StackName: ${self:custom.prefix}-s3
+  Resources:
+    FinetuneDataBucket:
+      Type: AWS::S3::Bucket
+      Properties:
+        BucketName: ${self:custom.prefix}-finetune-data-set
+  Outputs:
+    FinetuneDataBucketName:
+      Value:
+        Ref: FinetuneDataBucket
+      Export:
+        Name: ${self:custom.prefix}-finetune-data-bucket-name
diff --git a/serverless.yml b/serverless.yml
@@ -9,6 +9,8 @@ provider:
   lambdaHashingVersion: 20201221
   memorySize: 256
   timeout: 30
+  iam:
+    role: lambdaRole
 
 variablesResolutionMode: 20210326
 useDotenv: true
@@ -20,9 +22,16 @@ functions:
   - ${file(./resources/authorizer_function.yml)}
   - ${file(./resources/openai_functions.yml)}
 
+resources:
+  - ${file(./resources/iam.yml)}
+
 custom:
   prefix: ${self:service}-${self:provider.stage}
 
+  additionalStacks:
+    dynamodb: ${file(./resources/dynamodb.yml):stack}
+    s3: ${file(./resources/s3.yml):stack}
+
   pythonRequirements:
     dockerizePip: non-linux
     usePoetry: true
@@ -36,3 +45,4 @@ plugins:
   - serverless-python-requirements
   - serverless-offline
   - serverless-domain-manager
+  - serverless-plugin-additional-stacks