tuning-halving: AWS Lambda Deployment + Test

Signed-off-by: Alan Nair <[email protected]>

.github/workflows/e2e-tuning-halving.yml

@@ -70,11 +70,23 @@ jobs:
       id: buildx
       uses: docker/setup-buildx-action@v2
 
+    - name: Install AWS CLI
+      uses: unfor19/install-aws-cli-action@master
+      with:
+        version: '2'
+
+    - name: Set up Python Dependencies for AWS Jobs
+      working-directory: runner/aws_lambda_scripts
+      run: |
+        python3 -m pip install --upgrade pip
+        python3 -m pip install wheel ez_setup setuptools
+        python3 -m pip install -r requirements.txt
+
     - name: Build and Push
       working-directory: benchmarks/tuning-halving
       env:
         GOPRIVATE_KEY: ${{ secrets.XDT_REPO_ACCESS_KEY }}
-      run: make all-image-push
+      run: make all-push
 
   test-compose:
     name: Test Docker Compose
@@ -186,3 +198,29 @@ jobs:
         run: |
           kubectl delete -f ./service-driver.yaml --namespace default --wait
           kubectl delete -f ./service-trainer.yaml --namespace default --wait
+
+  test-aws-lambda:
+    name: Test AWS Lambda Deployment
+    needs: build-and-push
+    runs-on: ubuntu-20.04
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v3
+        with:
+          lfs: 'true'
+
+      - name: Set up Python Dependencies for AWS Jobs
+        working-directory: runner/aws_lambda_scripts
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install wheel ez_setup setuptools
+          python3 -m pip install -r requirements.txt
+
+      - name: Deploy and Test functions from ECR container
+        working-directory: runner/aws_lambda_scripts
+        run: |
+          python aws_actions.py deploy_lambdafn_from_ecr -n tuning-halving-driver -f tuning-halving-driver -p invoke_function,access_s3 -e '{"IS_LAMBDA":"true","TRAINER_FUNCTION":"tuning-halving-trainer","BUCKET_NAME":"vhive-stacking"}'
+          python aws_actions.py deploy_lambdafn_from_ecr -n tuning-halving-trainer -f tuning-halving-trainer -p access_s3 -e '{"IS_LAMBDA":"true","BUCKET_NAME":"vhive-stacking"}'
+          python aws_actions.py invoke_lambdafn -f stacking-training-driver

benchmarks/tuning-halving/Makefile

@@ -32,12 +32,18 @@ MODE = --load
 
 .PHONY: proto
 
-all: all-image
+all: all-image all-image-lambda
 
 all-image: driver-image trainer-image
 
+all-image-lambda: driver-image-lambda trainer-image-lambda
+
+all-push: all-image-push all-image-lambda-push
+
 all-image-push: driver-image-push trainer-image-push
 
+all-image-lambda-push: driver-image-lambda-push trainer-image-lambda-push
+
 driver-image: docker/Dockerfile proto/tuning_pb2_grpc.py proto/tuning_pb2.py driver/main.py
 	docker buildx build $(PLATFORM) \
 	-t vhiveease/tuning-halving-driver:latest \
@@ -54,12 +60,40 @@ trainer-image: docker/Dockerfile proto/tuning_pb2_grpc.py proto/tuning_pb2.py tr
 	-f docker/Dockerfile \
 	$(ROOT) $(MODE)
 
+driver-image-lambda: docker/Dockerfile.Lambda driver/main.py
+	docker buildx build $(PLATFORM) \
+	-t $(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com/tuning-halving-driver:latest \
+	--build-arg target_arg=driver \
+	-f docker/Dockerfile.Lambda \
+	$(ROOT) $(MODE)
+
+trainer-image-lambda: docker/Dockerfile.Lambda trainer/main.py
+	docker buildx build $(PLATFORM) \
+	-t $(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com/tuning-halving-trainer:latest \
+	--build-arg target_arg=trainer \
+	-f docker/Dockerfile.Lambda \
+	$(ROOT) $(MODE)
+
 driver-image-push: driver-image
 	docker push vhiveease/tuning-halving-driver:latest
 
 trainer-image-push: trainer-image
 	docker push vhiveease/tuning-halving-trainer:latest
 
+driver-image-lambda-push: driver-image-lambda
+	aws ecr get-login-password --region $(AWS_REGION) | \
+		docker login --username AWS --password-stdin \
+		$(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com
+	python $(ROOT)/runner/aws_lambda_scripts/aws_actions.py create_ecr_repo -n tuning-halving-driver
+	docker push $(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com/tuning-halving-driver:latest
+
+trainer-image-lambda-push: trainer-image-lambda
+	aws ecr get-login-password --region $(AWS_REGION) | \
+		docker login --username AWS --password-stdin \
+		$(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com
+	python $(ROOT)/runner/aws_lambda_scripts/aws_actions.py create_ecr_repo -n tuning-halving-trainer
+	docker push $(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com/tuning-halving-trainer:latest
+
 proto: proto/tuning.proto proto/helloworld.proto
 	python -m grpc_tools.protoc -I./proto --python_out=./proto --grpc_python_out=./proto ./proto/tuning.proto
 	python -m grpc_tools.protoc -I./proto --python_out=./proto --grpc_python_out=./proto ./proto/helloworld.proto

benchmarks/tuning-halving/docker/Dockerfile.Lambda

@@ -0,0 +1,37 @@
+# MIT License
+#
+# Copyright (c) 2022 Alan Nair and The vHive Ecosystem
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+FROM public.ecr.aws/lambda/python:3.9
+ARG target_arg
+ENV target=$target_arg
+
+# Copy function code
+COPY benchmarks/tuning-halving/${target}/* ${LAMBDA_TASK_ROOT}
+COPY utils/tracing/python/tracing.py ${LAMBDA_TASK_ROOT}
+COPY utils/storage/python/storage.py ${LAMBDA_TASK_ROOT}
+
+# Install the function dependencies
+COPY benchmarks/tuning-halving/requirements/aws_lambda.txt requirements.txt
+RUN  pip3 install -r requirements.txt --target "${LAMBDA_TASK_ROOT}"
+
+# Set the CMD to handler
+CMD [ "main.lambda_handler" ]

benchmarks/tuning-halving/driver/main.py

@@ -33,6 +33,7 @@
 
 if LAMBDA:
     import boto3
+    import json
 
 if not LAMBDA:
     import helloworld_pb2_grpc
@@ -102,6 +103,39 @@ def SayHello(self, request, context):
         self.driver.drive({'trainerfn': self.train})
         return helloworld_pb2.HelloReply(message=self.driver.storageBackend.bucket)
 
+class AWSLambdaDriver:
+    def __init__(self, XDTconfig=None):
+        self.driver = Driver()
+
+    def train(self, arg: dict) -> dict:
+        log.info("Invoke Trainer")
+        trainerArgs = {
+            'dataset_key': 'dataset_key',
+            'model_config': json.dumps(arg['model_config']),
+            'count': arg['count'],
+            'sample_rate': str(arg['sample_rate'])
+        }
+        resp = boto3.client("lambda").invoke(
+            FunctionName = os.environ.get('TRAINER_FUNCTION'),
+            InvocationType = 'RequestResponse',
+            LogType = 'None',
+            Payload = json.dumps(trainerArgs),
+        )
+        payloadBytes = resp['Payload'].read()
+        payloadJson = json.loads(payloadBytes)
+
+        return {
+            'model_key': payloadJson['model_key'],
+            'pred_key': payloadJson['pred_key'],
+            'score': float(payloadJson['score']),
+            'params': json.loads(payloadJson['params'])
+        }
+
+    def SayHello(self, event, context):
+        log.info("Driver received a request")
+        self.driver.drive({'trainerfn': self.train})
+        return self.driver.storageBackend.bucket
+
 def serve():
     transferType = os.getenv('TRANSFER_TYPE', S3)
     if transferType == S3:
@@ -124,6 +158,10 @@ def serve():
     else:
         log.fatal("Invalid Transfer type")
 
-if __name__ == '__main__':
+def lambda_handler(event, context):
+    awsLambdaDriver = AWSLambdaDriver()
+    return awsLambdaDriver.SayHello(event, context)
+
+if not LAMBDA and __name__ == '__main__':
     log.basicConfig(level=log.INFO)
     serve()

benchmarks/tuning-halving/requirements/aws_lambda.txt

@@ -0,0 +1,12 @@
+opentelemetry-api==1.3.0
+opentelemetry-exporter-zipkin==1.3.0
+opentelemetry-exporter-zipkin-json==1.3.0
+opentelemetry-exporter-zipkin-proto-http==1.3.0
+opentelemetry-instrumentation==0.22b0
+opentelemetry-instrumentation-grpc==0.22b0
+opentelemetry-sdk==1.3.0
+opentelemetry-semantic-conventions==0.22b0
+boto3==1.18.1
+environs==9.3.2
+scikit-learn==0.24.2
+numpy==1.20.1

benchmarks/tuning-halving/trainer/main.py

@@ -31,6 +31,9 @@
 
 LAMBDA = os.environ.get('IS_LAMBDA', 'no').lower() in ['true', 'yes', '1']
 
+if LAMBDA:
+    import json
+
 if not LAMBDA:
     import grpc
     import argparse
@@ -78,9 +81,27 @@ def Train(self, request, context):
             model_key=response['model_key'],
             pred_key=response['pred_key'],
             score=response['score'],
-            params=pickle.dumps(trainerArgs['model_config']),
+            params=request.model_config,
         )
 
+class AwsLambdaTrainer:
+    def __init__(self, XDTconfig=None):
+        self.trainer = Trainer(XDTconfig)
+
+    def Train(self, event, context):
+        trainerArgs = {
+            'dataset_key': event['dataset_key'],
+            'model_config': json.loads(event['model_config']),
+            'sample_rate': float(event['sample_rate']),
+            'count': event['count']
+        }
+        response = self.trainer.train(trainerArgs)
+        return {
+            'model_key': response['model_key'],
+            'pred_key': response['pred_key'],
+            'score': str(response['score']),
+            'params': event['model_config']
+        }
 
 def serve():
     transferType = os.getenv('TRANSFER_TYPE', S3)
@@ -99,7 +120,10 @@ def serve():
     else:
         log.fatal("Invalid Transfer type")
 
+def lambda_handler(event, context):
+    awsLambdaTrainer = AwsLambdaTrainer()
+    return awsLambdaTrainer.SayHello(event, context)
 
-if __name__ == '__main__':
+if not LAMBDA and __name__ == '__main__':
     log.basicConfig(level=log.INFO)
     serve()