template.yml

AWSTemplateFormatVersion: 2010-09-09
Description: "Foundation model benchmarking tool"

Parameters:
  S3BucketNameForRead:
    Default: sagemaker-fmbench-read
    Type: String
    Description: Name of the Amazon S3 bucket for holding datasets, scripts and tokenizer files. AWS region and account id would be suffixed automatically for uniqueness.
    MinLength: 1
    MaxLength: 63
    AllowedPattern: (?!(^xn--|.+-s3alias$))^[a-z0-9][a-z0-9-]{1,61}[a-z0-9]$
  S3BucketNameForWrite:
    Default: sagemaker-fmbench-write
    Type: String
    Description: Name of the Amazon S3 bucket for holding metrics and reports. AWS region and account id would be suffixed automatically for uniqueness.
    MinLength: 1
    MaxLength: 63
    AllowedPattern: (?!(^xn--|.+-s3alias$))^[a-z0-9][a-z0-9-]{1,61}[a-z0-9]$

Conditions:
  IsGovCloudPartition: !Equals
    - !Select [0, !Split [ "aws-us-gov", !Ref "AWS::Partition"]]
    - ""
  IsNotGovCloud: !Not [ !Condition IsGovCloudPartition ]

Resources:
  CodeRepository:
    Type: AWS::SageMaker::CodeRepository
    Condition: IsNotGovCloud
    Properties:
      GitConfig:
        RepositoryUrl: https://github.com/aws-samples/foundation-model-benchmarking-tool

  NotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: !Sub ${AWS::StackName}-notebook
      InstanceType: ml.m6i.xlarge
      RoleArn: !GetAtt NotebookRole.Arn
      DefaultCodeRepository: !If [ IsNotGovCloud, !GetAtt CodeRepository.CodeRepositoryName, !Ref "AWS::NoValue" ]
      VolumeSizeInGB: 200

  NotebookRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${AWS::StackName}-${AWS::Region}-role
      ManagedPolicyArns:
        - !Sub arn:${AWS::Partition}:iam::aws:policy/AmazonSageMakerFullAccess
        - !Sub arn:${AWS::Partition}:iam::aws:policy/AmazonS3FullAccess
        - !Sub arn:${AWS::Partition}:iam::aws:policy/AWSCloudFormationReadOnlyAccess
        - !Sub arn:${AWS::Partition}:iam::aws:policy/AmazonBedrockFullAccess
        - !Sub arn:${AWS::Partition}:iam::aws:policy/AWSPriceListServiceFullAccess 
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - "sts:AssumeRole"
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - "sts:AssumeRole"

  S3BucketForRead:
    Type: AWS::S3::Bucket
    Description: Amazon S3 bucket to hold source data
    Properties:
      BucketName: !Join
        - "-"
        - - !Ref S3BucketNameForRead
          - !Sub ${AWS::Region}
          - !Sub ${AWS::AccountId}

  S3BucketForWrite:
    Type: AWS::S3::Bucket
    Description: Amazon S3 bucket to hold metrics and reports
    Properties:
      BucketName: !Join
        - "-"
        - - !Ref S3BucketNameForWrite
          - !Sub ${AWS::Region}
          - !Sub ${AWS::AccountId}

  cleanupReadBucketOnDelete:
    Type: Custom::cleanupbucket
    Properties:
      ServiceToken: !GetAtt "DeleteS3Bucket.Arn"
      BucketName: !Ref S3BucketForRead
    DependsOn: S3BucketForRead

  cleanupWriteBucketOnDelete:
    Type: Custom::cleanupbucket
    Properties:
      ServiceToken: !GetAtt "DeleteS3Bucket.Arn"
      BucketName: !Ref S3BucketForWrite
    DependsOn: S3BucketForWrite

  DeleteS3Bucket:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Description: "Delete all objects in S3 bucket"
      Timeout: 300
      Role: !GetAtt "LambdaBasicExecutionRole.Arn"
      Runtime: python3.9
      Code:
        ZipFile: |
          import json, boto3, logging
          import cfnresponse
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)

          def lambda_handler(event, context):
              logger.info("event: {}".format(event))
              try:
                  bucket = event['ResourceProperties']['BucketName']
                  logger.info("bucket: {}, event['RequestType']: {}".format(bucket,event['RequestType']))
                  if event['RequestType'] == 'Delete':
                      s3 = boto3.resource('s3')
                      bucket = s3.Bucket(bucket)
                      for obj in bucket.objects.filter():
                          logger.info("delete obj: {}".format(obj))
                          s3.Object(bucket.name, obj.key).delete()

                  sendResponseCfn(event, context, cfnresponse.SUCCESS)
              except Exception as e:
                  logger.info("Exception: {}".format(e))
                  sendResponseCfn(event, context, cfnresponse.FAILED)

          def sendResponseCfn(event, context, responseStatus):
              responseData = {}
              responseData['Data'] = {}
              cfnresponse.send(event, context, responseStatus, responseData, "CustomResourcePhysicalID")

  CustomSGResource:
    Type: AWS::CloudFormation::CustomResource
    Properties:
      ServiceToken: !GetAtt "CustomFunctionCopyContentsToS3Bucket.Arn"

  LambdaBasicExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Path: /
      Policies:
        - PolicyName: S3Access
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub arn:${AWS::Partition}:logs:*:*:*
              - Effect: Allow
                Action:
                  - s3:*
                Resource: "*"

  CustomFunctionCopyContentsToS3Bucket:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Description: "Copies files from the Blog bucket to bucket in this account"
      MemorySize: 1024
      EphemeralStorage:
        Size: 2048
      Timeout: 900
      Role: !GetAtt "LambdaBasicExecutionRole.Arn"
      Runtime: python3.9
      Environment:
        Variables:
          READ_BUCKET: !Ref S3BucketForRead
          WRITE_BUCKET: !Ref S3BucketForWrite
          MY_AWS_REGION: !Ref AWS::Region
          ROLE_ARN: !GetAtt "NotebookRole.Arn"
      Code:
        ZipFile: |
          import os
          import boto3
          import logging
          import urllib3
          from urllib3.exceptions import HTTPError
          import cfnresponse

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          BLOGS_BUCKET = "aws-blogs-artifacts-public"
          SRC_PREFIX = "artifacts/ML-FMBT"
          MANIFEST = os.path.join(SRC_PREFIX, "manifest.txt")

          MANIFEST_URL = f"https://{BLOGS_BUCKET}.s3.amazonaws.com/{SRC_PREFIX}/manifest.txt"

          s3 = boto3.client("s3")
          http = urllib3.PoolManager()

          # s3://aws-blogs-artifacts-public/artifacts/ML-15729/docs/manifest.txt
          def lambda_handler(event, context):
              logger.info("got event {}".format(event))
              if event["RequestType"] == "Delete":
                  logger.info(
                      f"copy files function called at the time of stack deletion, skipping"
                  )
                  response = dict(files_copied=0, error=None)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, response)
                  return
                  # End DELETE

              try:
                  read_bucket = os.environ.get("READ_BUCKET")
                  write_bucket = os.environ.get("WRITE_BUCKET")

                  logger.info(f"Read Bucket: {read_bucket}")
                  logger.info(f"Write Bucket: {write_bucket}")

                  # obj = s3.get_object(Bucket=BLOGS_BUCKET, Key=MANIFEST)
                  manifest_data = []
                  try:
                      logger.debug(f"about to fetch manifest file: {MANIFEST_URL}")
                      response = http.request("GET", MANIFEST_URL)
                      manifest_data = response.data.decode("utf-8").splitlines()
                      manifest_data.append("manifest.txt")
                      logger.info(f"Fetched manifest file: {MANIFEST_URL}")
                      logger.debug(manifest_data)
                  except HTTPError as e:
                      logger.warning(f"Error downloading manifest {MANIFEST_URL}: {e}")
                      raise e
                  finally:
                      response.close()

                  ctr = 0
                  for fname in manifest_data:
                      is_config = fname.startswith("configs")
                      key = os.path.join(SRC_PREFIX, fname)
                      logger.info(f"going to read {key} from bucket={BLOGS_BUCKET}")

                      content = None
                      try:
                          file_url = f"https://{BLOGS_BUCKET}.s3.amazonaws.com/{key}"
                          logger.debug(f"about to fetch: {file_url}")
                          response = http.request("GET", file_url)
                          content = response.data.decode("utf-8")
                          logger.info(f"fetched: {file_url}")
                      except HTTPError as e:
                          logger.error(f"Error downloading {file_url}: {e}")
                          raise e
                      finally:
                          response.close()

                      # Retrieve the object from S3
                      if is_config:
                          logger.debug(f"found config file... updating content template: {fname}")
                          content = content.format(
                              region=os.environ.get("MY_AWS_REGION"),
                              role_arn=os.environ.get("ROLE_ARN"),
                              write_bucket=write_bucket,
                              read_bucket=read_bucket,
                              write_tmpdir="{write_tmpdir}",
                              read_tmpdir="{read_tmpdir}",
                          )
                          logger.info(f"Updating config file: {fname}")

                      s3.put_object(Bucket=read_bucket, Key=fname, Body=content)
                      logger.info(f"saving file to destination: Bucket:{read_bucket}/{key}")
                      ctr += 1

                  logger.info("Done!")
                  response = dict(files_copied=ctr, error=None)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, response)
              except Exception as e:
                  logger.error(e)
                  response = dict(files_copied=0, error=str(e))
                  cfnresponse.send(event, context, cfnresponse.FAILED, response)

              return

Outputs:
  S3BucketForRead:
    Value: !GetAtt S3BucketForRead.Arn
  S3BucketForWrite:
    Value: !GetAtt S3BucketForWrite.Arn
  FilesCopied:
    Description: Files copied
    Value: !GetAtt "CustomSGResource.files_copied"
  FileCopyError:
    Description: Files copy error
    Value: !GetAtt "CustomSGResource.error"
  Region:
    Description: Deployed Region
    Value: !Ref AWS::Region