[AI] FastAPI를 활용한 Tensorflow 모델 서빙

🔎 현재 상황

AI 기반의 필적 감정 서비스를 개발하게 되었다. 검사하려는 검증물과 비교 대상인 대조물 사이의 유사도 및 필적 특징을 추출하여 제공하는 서비스이다. 내가 담당하게 된 파트는 AI 서빙이었다. 평소 AI 기반 서비스가 어떻게 이루어지는지 많이 궁금하였고 이 참에 공부해보는 것도 좋겠다고 생각했다.

 

🤔 기술적 고민

현재 제공하는 AI 서빙 방법은 총 3가지이다.

  1. 웹 서버 프레임워크로 직접 서빙하기
  2. TensorFlow Serving 이용하기
  3. NVIDIA Triton

카카오 테크 블로그에서 비교된 아티클을 발견하여 참고하였다.

https://tech.kakaopay.com/post/model-serving-framework/

“FastAPI와 TensorFlowServing, Trioton과의 가장 큰 차이점은  Dynamic Batch Inference 지원 여부입니다. 해당 기능은 모델 서빙 속도에 가장 큰 영향을 미치는 기능 중 하나이지만, FastAPI의 경우 기본적으로 지원되지 않는 부분입니다.”

Dynamic Batch Inference는 여러 요청을 실시간으로 모아서 배치 처리해 성능을 높이는 기술이다. 해당 기술이 적용되니 않아 대규모 요청에 대해 성능이 떨어질 수 있다. 그러나 스트레스 테스트에서도 모두 준수한 성능을 나타내기 때문에 어떤 것을 선택하든지 성능적 큰 차이가 없다고 판단했다.

나는 결론적으로 FastAPI를 선택하기로 결정했다.

  1. S3 버킷의 다운로드 API가 필요하다.
  2. 코드가 간결하여 러닝커브가 낮다.
  3. 손쉬운 Swagger이 가능하다.

 

🚀 결정

Spring Boot와 FastAPI와의 REST로 통신할 수 있도록 구현했다.

 

🗂️ 고민한 점

AI 서빙 속도 최적화

처음 AI 모델 연산이 11초가 걸려 어떻게든 속도 최적화가 필요한 상태였다. 그러나 AI 모델의 연산 자체가 11초가 걸리기 때문에 서버 상에서 이를 줄일 수 있는 방법이 있을까 고민했다. 다행히 이미지를 한줄씩 대조하는 로직을 하나의 이미지 통채로 대조하는 로직으로 변경 후 3초로 줄일 수 있었지만, 추가적인 해결책이 있지 않을까 생각했다. 캐싱을 통해 중복된 요청에 대한 연산결과를 저장할까 고민했지만, 서비스 특성상 중복 요청이 존재하기 어렵다고 판단하여 일단 보류하기로 했다.

 

AI 모델과 서버의 분리

함께 협업하는 AI 개발자가 웹 서버에 대한 지식이 없어 내가 소스코드를 바탕으로 서빙하기로 결정했다. 만약 AI 소스코드가 변경될 경우, AI 개발자 분이 직접 서버에 반영할 수 있도록 AI 모델과 서버를 분리하려고 했다. Tensorflow Serving을 통해 웹 서버와 분리하여 서빙이 가능하지만, S3 이용이 어렵다는 문제점이 있다. 어쩔 수 없이 FastAPI를 통해 AI 모델과 통합 운영하는 방안으로 결정했다.

다만 AI 개발자가 쉽게 소스코드를 변경할 수 있도록 AI 패키지와 App 패키지를 구분하여 구현했다.

 

도커를 사용하지 않은 이유

도커를 통해 AI 모델을 서빙할 경우 매우 간편하다. 그러나 이번에는 이용하지 않기로 했다. EC2 프리티어를 사용하기 위해 t2.micro 인스턴스에 서버를 운영 중이다. 따라서 한정된 메모리와 램을 사용한다. 이번 서비스의 경우, 실시간 분석 기능이 포함되어 있어 CPU 및 메모리 사용량이 많을 것으로 예상된다. 조금이나마 자원 소모를 줄이기 위해 Docker의 컨테이너 레이어 사용없이 OS 단에서 실행될 수 있도록 구현했다.

 

⏰ 시퀀스 다이어그램

  1. 사용자가 분석할 검증물을 Presigned URL을 통해 S3에 업로드한다. (앱 -> 스프링)

1.1. 업로드된 S3의 검증물 URL을 반환한다. (스프링 -> 앱)

  1. 사용자가 분석할 대조물을 Presigned URL을 통해 S3에 업로드한다. (앱 -> 스프링)

2.1 업로드된 S3의 대조물 URL을 반환한다. (스프링 -> 앱)

  1. 감정하기를 의뢰한다. (앱 -> 스프링)
  2. AI 서버에게 감정하기를 의뢰한다. (스프링 -> FastAPI)

4.1 감정된 결과를 JSON으로 응답한다. (FastAPI -> 스프링)

4.2 Result 객체를 만들어 DB에 저장한다. (스프링)

  1. 감정 결과를 사용자에게 최종 전달한다. (스프링 -> 앱)

 

🤯 발생할 수 있는 예상 문제

스프링 서버에서 외부 API를 호출하기 때문에 네트워크 I/O가 발생한다. 거기에 더해 실시간으로 aI 모델을 돌리는 만큼 더 긴 네트워크 I/O가 발생하여 대규모 트래픽이 올 경우, 병목이 발생할 수 있다고 판단했다.

이를 해결하기 위해 비동기 요청을 통해 요청 간의 병목을 어느정도 해결할 수 있다고 판단했다. 또한 대규모 트래픽으로 발생하는 서버 부하 문제는 스케일 아웃과 Circuit Breaker 패턴을 적용하여 서버의 가용성을 향상시킬 수 있겠다고 생각했다.

비동기 적용에 대해선 해당 이슈를 참고하면된다!

https://github.com/dog-feet-bird-feet/server/issues/40

 

스케일 아웃을 하지 않은 이유

K6를 통해 성능을 측정해본 결과 10명 사용자가 동시에 요청을 보낼 때 병목이 발생했다. 최대 응답 시간은 10초, 평균 5초의 응답 요청이 발생했다. 그러나 이미지를 분석하는 서비스의 특성상 빠른 응답은 사용자에게 신뢰성을 떨어뜨리는 요소라고 판단했다. 만약 분석 툴을 사용하여 결과가 1~2초 안에 나올 경우, 제대로 측정이 된 것인지 의문을 가질 수 있다고 생각했다.

또한, 현재 사용 중인 t2.micro 환경에서 CPU 사용률이 50% 미만으로 크게 나오지 않아 스케일 아웃 적용도 급하지 않다고 판단했다. 아마 이미지 다운로드 및 삭제하는 디스크 I/O로 인해 발생한 병목이라고 생각된다.

 

🔨 구현

Spring Boot

@PostMapping("/appraisal")
public ResponseEntity<AppraisalResponse> appraisal(@AuthenticationPrincipal PrincipalDetails principalDetails, @RequestBody AppraisalRequest request) {
    AppraisalResponse response = appraisalService.appraise(principalDetails.getAuthenticatedMember().getMemberId(), request);
    return ResponseEntity.ok(response);
}

/appraisal로 요청을 보내면 JWT 토큰을 통해 받은 memberId와 업로드된 검증물과 대조물 이미지 URL을 받아 Service 레이어로 넘긴다.

public AppraisalResponse appraise(final long memberId, final AppraisalRequest request) {
    RestClient restClient = RestClient.builder()
            .baseUrl(fastApiEndpoint)
            .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
            .build();

    AppraisalAIResponse response = restClient
            .post()
            .uri(uriBuilder -> uriBuilder.path("/analyze").build())
            .body(request)
            .retrieve()
            .body(AppraisalAIResponse.class);

    return AppraisalResponse.from(saveAppraisal(memberId, Objects.requireNonNull(response)));
}

public AnalysisResult saveAppraisal(long memberId, AppraisalAIResponse response) {
    Member member = memberRepository.findById(memberId).orElseThrow(() -> BaseException.from(MEMBER_NOT_FOUND));
    AnalysisResult result = response.toEntity();
    result.update(member);
    resultRepository.save(result);
    return result;
}

RestClient를 통해 AI 서버와 통신하도록 구현을 했고, 유저에게 응답하기 전 데이터베이스에 저장하는 코드를 삽입하였다. 그리고 이 과정에 대해 트랜잭션을 적용하였다.

 

AI 서버 (FastAPI)

@router.post("/analyze")
def generate(request: Request, req_body: AnalyzeRequest):
    model = request.app.state.model

    # 1. 이미지 다운로드
    reference_hash = s3_manager.download_reference_images(req_body.comparisonImageUrls)
    test_hash = s3_manager.download_test_image(req_body.verificationImageUrl)

    # 2. 감정 시작
    appraisal_response = ai_model.analyze(model, reference_hash, test_hash)
    appraisal_response.verificationImageUrl = req_body.verificationImageUrl

    # 3. 이미지 삭제
    await s3_manager.delete_reference_images(reference_hash)
    await s3_manager.delete_test_image(test_hash)

    return appraisal_response

s3_manager를 통해 S3에 업로드된 이미지를 다운로드 및 삭제하도록 구현했고, ai_model의 analyze 메서드를 호출하여 이미지 분석이 되도록 구현했다. 그리고 분석된 결과를 response 형식으로 저장하여 응답한다.

더 자세한 코드는 깃허브를 참고해주세요!

 

Spring: https://github.com/dog-feet-bird-feet/server

FastAPI: https://github.com/dog-feet-bird-feet/ai

 

✅ 결과

현재 로컬에서 돌린 결과 576ms로 매우 빠른 성능을 자랑한다. 그러나 현재 배포된 EC2 환경의 경우 싱글 코어이기 때문에 평균 3ms 이상의 성능이 나오는 것을 확인했다.