AWS Lambda环境下导入tensorflow.keras耗时长的原因及解决方法

我尝试在AWS Lambda中加载并使用keras模型，但在Lambda函数中从tensorflow导入keras的耗时非常久。但奇怪的是，该操作在SageMaker中运行时耗时很短。我想了解出现该现象的原因，以及对应的解决方法。

我使用Serverless Framework部署函数，相关的处理函数代码与serverless.yml配置如下。我使用挂载了EFS卷的EC2实例安装所有依赖并存储在该EFS卷中，通过pip的-t参数将依赖安装到EFS路径下，例如安装tensorflow的命令如下：

sudo miniconda3/envs/devenv/bin/pip install tensorflow -t /mnt/efs/fs1/lib

其中/mnt/efs/fs1/lib是EFS上存储依赖的目录，模型文件存储在S3中。

我先在SageMaker笔记本中使用以下代码完成了模型加载逻辑的原型验证：

import time
start = time.time()

from tensorflow import keras
print('keras: {}'.format(time.time()-start))
import boto3
import os
import zipfile
print('imports: {}'.format(time.time()-start))

modelPath = '***model.zip'
bucket = 'predictionresources'


def load_motion_model():
    s3 = boto3.client('s3')
    s3.download_file(bucket, modelPath, 'model.motionmodel.zip')
    
    with zipfile.ZipFile('model.motionmodel.zip', 'r') as zip_ref:
        zip_ref.extractall('model.motionmodel')
        
    return keras.models.load_model('model.motionmodel/'+os.listdir('model.motionmodel')[0])

model = load_motion_model()

print('total time: {}'.format(time.time()-start))

上述代码的运行输出结果如下：

keras: 2.0228586196899414
imports: 2.0231151580810547
total time: 3.0635251998901367

即包含所有导入操作在内，该逻辑总执行耗时约3秒。

但当我通过Serverless Framework将函数部署到AWS Lambda后，keras的导入耗时大幅上升，对应的Lambda函数代码（与原型逻辑一致，仅增加了Lambda处理函数的封装）与serverless.yml配置如下：

处理函数（Handler）代码 #

try:
    import sys
    import os
    sys.path.append(os.environ['MNT_DIR']+'/lib0')  # nopep8 # noqa
except ImportError:
    pass

#returns the version of all dependancies
def test(event, context):
    print('TEST LOADING')

    import time
    start = time.time()

    from tensorflow import keras
    print('keras: {}'.format(time.time()-start))
    import boto3
    import os
    import zipfile
    print('imports: {}'.format(time.time()-start))

    modelPath = '**********nmodel.zip'
    bucket = '***********'


    def load_motion_model():
        s3 = boto3.client('s3')
        s3.download_file(bucket, modelPath, 'model.motionmodel.zip')
        
        with zipfile.ZipFile('model.motionmodel.zip', 'r') as zip_ref:
            zip_ref.extractall('model.motionmodel')
            
        return keras.models.load_model('model.motionmodel/'+os.listdir('model.motionmodel')[0])

    model = load_motion_model()

    print('total time: {}'.format(time.time()-start))

    body = {
        'message': 'done!'
    }

    response = {
        "statusCode": 200,
        "body": json.dumps(body)
    }

    return response

（注：我已知该代码会因Lambda临时目录无写权限运行失败，模型文件需要保存到/tmp/路径下）

Serverless.yml配置 #

service: test4KerasTest

plugins:
  - serverless-pseudo-parameters

custom:
  efsAccessPoint: fsap-*****
  LocalMountPath: /mnt/efs
  subnetsId: subnet-*****
  securityGroup: sg-*****

provider:
  name: aws
  runtime: python3.6
  region: us-east-2
  timeout: 30

  iamRoleStatements:
    - Effect: Allow
      Action:
        - s3:GetObject
        - s3:PutObject
      Resource: 'arn:aws:s3:::predictionresources/*'

package:
  exclude:
    - node_modules/**
    - .vscode/**
    - .serverless/**
    - .pytest_cache/**
    - __pychache__/**

functions:
  test:
    handler: handler.test
    environment: # Service wide environment variables
      MNT_DIR: ${self:custom.LocalMountPath}
      BUCKET: predictionresources
      REGION: us-east-2
    vpc:
      securityGroupIds:
        - ${self:custom.securityGroup}
      subnetIds:
        - ${self:custom.subnetsId}
    iamManagedPolicies:
      - arn:aws:iam::aws:policy/AmazonElasticFileSystemClientReadWriteAccess
      - arn:aws:iam::aws:policy/AmazonS3FullAccess
    events:
      - http:
          path: test
          method: get
    fileSystemConfig:
      localMountPath: '${self:custom.LocalMountPath}'
      arn: 'arn:aws:elasticfilesystem:${self:provider.region}:#{AWS::AccountId}:access-point/${self:custom.efsAccessPoint}'

部署后从CloudWatch观察到的运行结果显示：

可以看到，Lambda环境下keras的导入耗时远高于其他模块，即使调整不同模块的导入顺序，keras的导入耗时始终处于不合理的高水平。由于API Gateway的限制，该函数总运行时长不能超过30秒，因此需要找到合适的方法，缩短Lambda环境下导入keras的耗时。

内容的提问来源于stack exchange，提问作者Daniel Warfield