Cách AWS Lambda SnapStart loại bỏ khởi động nguội cho Serverless Machine Learning Inference

Nov 29 2022
Những thách thức khi bắt đầu suy luận ML Một trong những thách thức chính với Suy luận học máy không có máy chủ luôn là một khởi đầu lạnh lùng. Và trong trường hợp suy luận ML, có nhiều yếu tố góp phần vào nó: Tính năng SnapStart Với ​​tính năng SnapStart mới được công bố cho AWS Lambda, khởi động nguội được thay thế bằng SnapStart.

Những thách thức khi bắt đầu suy luận ML

Một trong những thách thức chính với Serverless Machine Learning Inference luôn là một khởi đầu khó khăn. Và trong trường hợp suy luận ML, có nhiều thứ góp phần vào nó:

  • khởi tạo thời gian chạy
  • tải thư viện và phụ thuộc
  • tự tải mô hình (từ S3 hoặc gói)
  • khởi tạo mô hình

Tính năng SnapStart

Với tính năng SnapStart mới được công bố dành cho AWS Lambda, khởi động nguội được thay thế bằng SnapStart. AWS Lambda sẽ tạo một ảnh chụp nhanh được mã hóa, không thay đổi về trạng thái bộ nhớ và ổ đĩa, đồng thời sẽ lưu vào bộ đệm ẩn để sử dụng lại. Ảnh chụp nhanh này sẽ có mô hình ML được tải trong bộ nhớ và sẵn sàng sử dụng.

Những điều cần lưu ý:

  • [Thời gian chạy Java] SnapStart hiện chỉ được hỗ trợ cho thời gian chạy Java. Điều đó làm tăng thêm các hạn chế, nhưng ONNX hoạt động trên Java và có thể chạy ONNX với SnapStart.
  • [Tải mô hình] Tải mô hình phải xảy ra trong bước khởi tạo, không phải trong bước chạy và mô hình phải được sử dụng lại giữa các lần chạy. Trong java, đó là một khối tĩnh. Điều tốt là chúng tôi không bị giới hạn bởi thời gian chờ chức năng để tải mô hình và thời gian khởi tạo tối đa là 15 phút.
  • [Snap-Resilient] SnapStart có những hạn chế cụ thể — tính duy nhất vì SnapStart sử dụng ảnh chụp nhanh. Ví dụ, điều đó có nghĩa là nếu một hạt giống ngẫu nhiên được xác định trong giai đoạn init thì tất cả lệnh gọi lambda sẽ có cùng một trình tạo. Đọc thêm về cách làm cho Lambda trở nên linh hoạt tại đây .

Một ví dụ với ONNX và SnapStart được cung cấp công khai tại đây và có thể được sử dụng với Sam để triển khai và kiểm tra điểm cuối ONNX Inception V3.

Để làm nổi bật kiến ​​trúc cho SnapStart trong trường hợp ONNX:

  • onnxSession — có mô hình được tải sẵn và được sử dụng lại giữa các lần gọi.
  • getOnnxSession — tải mô hình nếu mô hình chưa được tải trước đó và bỏ qua mô hình nếu mô hình đã được tải trước đó.
  • khối tĩnh — chạy mã trong quá trình tạo SnapStart. Đây là phần quan trọng — mã trong trình xử lý sẽ không được chạy trong quá trình tạo ảnh chụp nhanh.
  • package onnxsnapstart;
    
    /**
     * Handler for Onnx predictions on Lambda function.
     */
    public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
    
        // Onnx session with preloaded model which will be reused between invocations and will be
        // initialized as part of snapshot creation
        private static OrtSession onnxSession;
    
        // Returns Onnx session with preloaded model. Reuses existing session if exists.
        private static OrtSession getOnnxSession() {
            String modelPath = "inception_v3.onnx";
            if (onnxSession==null) {
              System.out.println("Start model load");
              try (OrtEnvironment env = OrtEnvironment.getEnvironment("createSessionFromPath");
                OrtSession.SessionOptions options = new SessionOptions()) {
              try {
                OrtSession session = env.createSession(modelPath, options);
                Map<String, NodeInfo> inputInfoList = session.getInputInfo();
                Map<String, NodeInfo> outputInfoList = session.getOutputInfo();
                System.out.println(inputInfoList);
                System.out.println(outputInfoList);
                onnxSession = session;
                return onnxSession;
              }
              catch(OrtException exc) {
                exc.printStackTrace();
              }
            }
            }
            return onnxSession;
        }
    
        // This code runs during snapshot initialization. In the normal lambda that would run in init phase.
        static {
            System.out.println("Start model init");
            getOnnxSession();
            System.out.println("Finished model init");
        }
    
        // Main handler for the Lambda
        public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
            Map<String, String> headers = new HashMap<>();
            headers.put("Content-Type", "application/json");
            headers.put("X-Custom-Header", "application/json");
    
    
            float[][][][] testData = new float[1][3][299][299];
    
            try (OrtEnvironment env = OrtEnvironment.getEnvironment("createSessionFromPath")) {
                OnnxTensor test = OnnxTensor.createTensor(env, testData);
                OrtSession session = getOnnxSession();
                String inputName = session.getInputNames().iterator().next();
                Result output = session.run(Collections.singletonMap(inputName, test));
                System.out.println(output);
            }
            catch(OrtException exc) {
                exc.printStackTrace();
            }
    
    
            APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent().withHeaders(headers);
            String output = String.format("{ \"message\": \"made prediction\" }");
    
            return response
                    .withStatusCode(200)
                    .withBody(output);
        }
    }
    

  • Lambda truyền thống
  • Picked up JAVA_TOOL_OPTIONS: -XX:+TieredCompilation -XX:TieredStopAtLevel=1
    Start model init
    Start model load
    {x.1=NodeInfo(name=x.1,info=TensorInfo(javaType=FLOAT,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,shape=[1, 3, 299, 299]))}
    {924=NodeInfo(name=924,info=TensorInfo(javaType=FLOAT,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,shape=[1, 1000]))}
    Finished model init
    START RequestId: 27b8e0e6-2e26-4356-a015-540fc79c080a Version: $LATEST
    ai.onnxruntime.OrtSession$Result@e580929
    END RequestId: 27b8e0e6-2e26-4356-a015-540fc79c080a
    REPORT RequestId: 27b8e0e6-2e26-4356-a015-540fc79c080a Duration: 244.99 ms Billed Duration: 245 ms Memory Size: 1769 MB Max Memory Used: 531 MB Init Duration: 8615.62 ms
    

    RESTORE_START Runtime Version: java:11.v15 Runtime Version ARN: arn:aws:lambda:us-east-1::runtime:0a25e3e7a1cc9ce404bc435eeb2ad358d8fa64338e618d0c224fe509403583ca
    RESTORE_REPORT Restore Duration: 571.67 ms
    START RequestId: 9eafdbf2-37f0-430d-930e-de9ca14ad029 Version: 1
    ai.onnxruntime.OrtSession$Result@47f6473
    END RequestId: 9eafdbf2-37f0-430d-930e-de9ca14ad029
    REPORT RequestId: 9eafdbf2-37f0-430d-930e-de9ca14ad029 Duration: 496.51 ms Billed Duration: 645 ms Memory Size: 1769 MB Max Memory Used: 342 MB Restore Duration: 571.67 ms
    

Chúng tôi vẫn có độ trễ bổ sung do khôi phục ảnh chụp nhanh, nhưng hiện tại phần đuôi của chúng tôi đã bị rút ngắn đáng kể và chúng tôi không có yêu cầu nào có thể mất hơn 2,5 giây.

  • Lambda truyền thống
  • Percentage of the requests served within a certain time (ms)
      50%    352
      66%    377
      75%    467
      80%    473
      90%    488
      95%   9719
      98%  10329
      99%  10419
     100%  12825
    

    50%    365
      66%    445
      75%    477
      80%    487
      90%    556
      95%   1392
      98%   2233
      99%   2319
     100%   2589 (longest request)