Feature/sagemaker integration

.gitignore

@@ -94,4 +94,5 @@ _build/
 # setuptools-scm/
 src/*/_version.py
 data/
-models/
+models/
+.aws-sam/

deployment/samconfig.yaml

@@ -0,0 +1,30 @@
+# More information about the configuration file can be found here:
+# https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-config.html
+version: 0.1
+
+default:
+  global:
+    parameters:
+      stack_name: sunscan-serverless-dev
+  build:
+    parameters:
+      cached: true
+      parallel: true
+  validate:
+    parameters:
+      lint: true
+  deploy:
+    parameters:
+      region: us-east-1
+      s3_bucket: project-deployment-templates
+      s3_prefix: sunscan-serverless
+      capabilities:
+        - CAPABILITY_NAMED_IAM
+        - CAPABILITY_AUTO_EXPAND
+      tags:
+        - Product=Sunscan
+        - Environment=dev
+      parameter_overrides:
+        - Mode=dev
+        - AttachmentsBucketName=infra-attachments-dev
+        - ArtifactsBucketName=infra-artifacts-dev

deployment/template.yaml

@@ -0,0 +1,188 @@
+AWSTemplateFormatVersion: 2010-09-09
+Transform: AWS::Serverless-2016-10-31
+Description: sunscan-serverless
+
+Parameters:
+  AttachmentsBucketName:
+    Type: String
+    Description: Attachments bucket name
+  ArtifactsBucketName:
+    Type: String
+    Description: Artifacts bucket name
+  Mode:
+    Type: String
+    Description: Environment context
+    AllowedValues:
+      - dev
+      - test
+      - prod
+
+
+Globals:
+  Function:
+    Timeout: 30
+    Runtime: python3.10
+    Environment:
+      Variables:
+        LOG_LEVEL: INFO
+        POWERTOOLS_LOGGER_SAMPLE_RATE: 0.1
+        POWERTOOLS_LOGGER_LOG_EVENT: true
+
+
+Resources:
+  LambdaRole:
+    Type: AWS::IAM::Role
+    Properties:
+      Path: /service-role/
+      RoleName: !Sub 'SunscanLambdaRole-${Mode}'
+      Description: Sunscan Lambda service role.
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: lambda.amazonaws.com
+            Action: sts:AssumeRole
+      ManagedPolicyArns:
+        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+
+  StatesRole:
+    Type: AWS::IAM::Role
+    Properties:
+      Path: /service-role/
+      RoleName: !Sub 'SunscanStatesRole-${Mode}'
+      Description: Sunscan States service role.
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: states.amazonaws.com
+            Action: sts:AssumeRole
+      Policies:
+        - PolicyName: sunscan-states-policy
+          PolicyDocument:
+            Version: '2012-10-17'
+            Statement:
+              - Effect: Allow
+                Action:
+                  - xray:PutTraceSegments
+                  - xray:PutTelemetryRecords
+                  - xray:GetSamplingRules
+                  - xray:GetSamplingTargets
+                Resource: '*' # NOSONAR
+              - Effect: Allow
+                Action:
+                  - events:PutTargets
+                  - events:PutRule
+                  - events:DescribeRule
+                Resource: !Sub 'arn:aws:events:${AWS::Region}:${AWS::AccountId}:rule/StepFunctionsGetEventsForSageMakerProcessingJobsRule'
+              - Effect: Allow
+                Action: lambda:InvokeFunction
+                Resource: !GetAtt BuildPayloadFunction.Arn
+              - Effect: Allow
+                Action:
+                  - sagemaker:CreateProcessingJob
+                  - sagemaker:DescribeProcessingJob
+                  - sagemaker:StopProcessingJob
+                  - sagemaker:AddTags
+                Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:processing-job/*'
+              - Effect: Allow
+                Action: iam:PassRole
+                Resource: !GetAtt SageMakerRole.Arn
+
+  SageMakerRole:
+    Type: AWS::IAM::Role
+    Properties:
+      Path: /service-role/
+      RoleName: !Sub 'SunscanSageMakerRole-${Mode}'
+      Description: Sunscan SageMaker service role.
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: sagemaker.amazonaws.com
+            Action: sts:AssumeRole
+      ManagedPolicyArns:
+        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+        - arn:aws:iam::aws:policy/AWSXrayWriteOnlyAccess
+      Policies:
+        - PolicyName: sunscan-lambda-policy
+          PolicyDocument:
+            Version: '2012-10-17'
+            Statement:
+              - Effect: Allow
+                Action:
+                  - cloudwatch:PutMetricData
+                  - logs:CreateLogStream
+                  - logs:CreateLogGroup
+                  - logs:DescribeLogStreams
+                  - logs:PutLogEvents
+                  - ecr:GetAuthorizationToken
+                Resource: '*' # NOSONAR
+              - Effect: Allow
+                Action:
+                  - ecr:BatchCheckLayerAvailability
+                  - ecr:GetDownloadUrlForLayer
+                  - ecr:BatchGetImage
+                Resource: !GetAtt RoofEnergyInferenceRepository.Arn
+              - Effect: Allow
+                Action: s3:ListBucket
+                Resource:
+                  - !Sub 'arn:aws:s3:::${AttachmentsBucketName}'
+                  - !Sub 'arn:aws:s3:::${ArtifactsBucketName}'
+              - Effect: Allow
+                Action:
+                  - s3:GetObject
+                  - s3:PutObject
+                Resource: !Sub 'arn:aws:s3:::${AttachmentsBucketName}/sunscan/*'
+              - Effect: Allow
+                Action: s3:GetObject
+                Resource: !Sub 'arn:aws:s3:::${ArtifactsBucketName}/sunscan/sagemaker/*'
+
+  RoofEnergyInferenceRepository:
+    Type: AWS::ECR::Repository
+    Properties:
+      ImageTagMutability: MUTABLE
+      RepositoryName: !Sub 'sunscan/roof-energy-inference-${Mode}'
+
+  BuildPayloadFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      Role: !GetAtt LambdaRole.Arn
+      FunctionName: !Sub 'sunscan-build-payload-${Mode}'
+      Description: Prepare input data to parameters on workflow.
+      CodeUri: ../src/lambda/build-payload
+      Handler: lambda_function.lambda_handler
+      Environment:
+        Variables:
+          ATTACHMENTS_BUCKET_NAME: !Ref AttachmentsBucketName
+          BASE_CODE_S3URI: !Sub 's3://${ArtifactsBucketName}/sunscan/sagemaker/roof-energy'
+
+  BuildPayloadFunctionLogGroup:
+    Type: AWS::Logs::LogGroup
+    Properties:
+      LogGroupName: !Sub '/aws/lambda/${BuildPayloadFunction}'
+      RetentionInDays: 30
+
+  RoofEnergyInferenceWorkflow:
+    Type: AWS::Serverless::StateMachine
+    Properties:
+      Role: !GetAtt StatesRole.Arn
+      Name: !Sub 'sunscan-roof-energy-inference-workflow-${Mode}'
+      Definition:
+        Fn::Transform:
+          Name: AWS::Include
+          Parameters:
+            Location: ../src/stepfunctions/RoofEnergyInferenceWorkflow.yaml
+      DefinitionSubstitutions:
+        Mode: !Ref Mode
+        BuildPayloadFunctionArn: !GetAtt BuildPayloadFunction.Arn
+        InferenceRoleArn: !GetAtt SageMakerRole.Arn
+        InferenceImageUri: !Sub '${RoofEnergyInferenceRepository.RepositoryUri}:2.1.0-gpu-py310'
+        InferenceInstanceType: ml.p3.2xlarge
+      Tags:
+        Application: Sunscan
+        Product: Infradigital
+        Environment: !Ref Mode

events/roof-energy-event.json

@@ -0,0 +1,27 @@
+{
+    "version": "0",
+    "id": "8728a076-d5d1-c624-124a-946b1b65e33b",
+    "detail-type": "Infradigital Application Task Requested",
+    "source": "infra.core",
+    "account": "123456789123",
+    "time": "2023-05-01T19:08:34Z",
+    "region": "us-east-1",
+    "resources": [
+        "arn:aws:states:us-east-1:123456789123:stateMachine:infra-task-manager-dev",
+        "arn:aws:states:us-east-1:123456789123:execution:infra-task-manager-dev:05cd1ed4-6277-383d-a3ef-e36e63862e57"
+    ],
+    "detail": {
+        "TaskKey": "a4066ee9-09e3-4f58-895a-6ef9f0445456",
+        "AppServiceSlug": "sunscan#roof_energy_inference",
+        "UserSub": "2ecf2bb8-c700-4073-9d48-2745815dcd0d",
+        "TenantId": "iadb:ine:tsp",
+        "AccessLevel": "user",
+        "Name": "Miami neighbourhood",
+        "Inputs": {
+            "BoundingBox": [ -73.2316, 9.587, -73.2309, 9.5875 ],
+            "PanelSize": 4,
+            "AvailableArea": 0.5,
+            "PanelPower": 400
+        }
+    }
+}

src/lambda/build-payload/inference_builder.py

@@ -0,0 +1,48 @@
+import parameter_mappers as mappers
+
+MODEL_NAME = 'sam_vit_h_4b8939.pth'
+
+
+class InferenceBuilder(object):
+    def __init__(self, task_key: str, base_code_s3uri: str, attachments_bucket_name: str):
+        self.task_key = task_key
+        self.base_code_s3uri = base_code_s3uri
+        self.attachments_bucket_name = attachments_bucket_name
+
+    def build_inference_parameters(self, payload):
+        return {
+            'JobName': mappers.get_unique_job_name(self.task_key, 'roof-energy'),
+            'ContainerArguments': payload['container_args'],
+            'InputConfig': [
+                mappers.get_processing_input(
+                    'code',
+                    self.base_code_s3uri + '/code',
+                    '/opt/ml/processing/code'
+                ),
+                mappers.get_processing_input(
+                    'models',
+                    self.base_code_s3uri + f'/{MODEL_NAME}',
+                    '/opt/ml/processing/models'
+                )
+            ],
+            'OutputConfig': [
+                mappers.get_processing_output('outputs',
+                    '/opt/ml/processing/outputs', payload['inference_outputs_uri'])
+            ]
+        }
+
+    def run(self, inputs: dict, user_sub: str):
+        inference_container_args = [
+            '--bounding-box', *[str(coord) for coord in inputs['BoundingBox']],
+            '--panel-size', str(inputs['PanelSize']),
+            '--available-area', str(inputs['AvailableArea']),
+            '--panel-power', str(inputs['PanelPower'])
+        ]
+
+        payload = {
+            'task_key': self.task_key,
+            'container_args': inference_container_args,
+            'inference_outputs_uri': f's3://{self.attachments_bucket_name}/sunscan/user:{user_sub}/roof_energy/{self.task_key}/outputs/'
+        }
+
+        return self.build_inference_parameters(payload)

src/lambda/build-payload/lambda_function.py

@@ -0,0 +1,26 @@
+import os
+import logging
+
+from inference_builder import InferenceBuilder
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+logging.getLogger('botocore').setLevel(logging.CRITICAL)
+
+BASE_CODE_S3URI = os.getenv('BASE_CODE_S3URI')
+ATTACHMENTS_BUCKET_NAME = os.getenv('ATTACHMENTS_BUCKET_NAME')
+
+
+def lambda_handler(event: dict, context):
+    builder = InferenceBuilder(
+        event['TaskKey'],
+        BASE_CODE_S3URI,
+        ATTACHMENTS_BUCKET_NAME
+    )
+
+    payload = builder.run(
+        event['Inputs'],
+        user_sub=event['UserSub']
+    )
+
+    return payload

src/lambda/build-payload/parameter_mappers.py

@@ -0,0 +1,33 @@
+from datetime import datetime, timezone
+
+def get_unique_job_name(uuid_key: str, *references: list):
+    """ Returns a unique job name based on a given uuid_key
+        and the current timestamp """
+
+    return '-'.join([
+        *uuid_key.split('-')[:2],
+        datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S'),
+        *references
+    ])[:63]
+
+def get_processing_input(name: str, inputs_s3_uri: str, output_local_path: str):
+    return {
+        'InputName': name,
+        'S3Input': {
+            'S3DataType': 'S3Prefix',
+            'S3Uri': inputs_s3_uri,
+            'LocalPath': output_local_path,
+            'S3DataDistributionType': 'FullyReplicated',
+            'S3InputMode': 'File'
+        }
+    }
+
+def get_processing_output(name: str, input_local_path: str, outputs_s3_uri: str):
+    return {
+        'OutputName': name,
+        'S3Output': {
+            'LocalPath': input_local_path,
+            'S3Uri': outputs_s3_uri,
+            'S3UploadMode': 'EndOfJob'
+        }
+    }