[Spring] AWS SageMaker 엔드포인트 연결 시도하기

문제 상황

기존 방식대로 Spring WebFlux 에서 HTTP Client로 사용되는 비동기적으로 작동하는 모듈인 WebClient를 활용하여 SageMaker 엔드포인트에 POST 요청을 시도했다.

 

그런데 CloudWatch에서 요청 시도에 대한 로깅은 찍히는데, 웹서버로 리턴이 들어오지 않는 상황이었다.

 

알고 보니, AWS SageMaker는 사용하는 IAM에 대한 정보를 바탕으로 AWS Signature를 Header에 담아 전송해야하는 것이었다.

 

이를 해결하기 위해 Spring에서는 aws sagemaker-runner 라이브러리를 사용해야 한다.

 

해결

 

1. build.gradle에 다음과 같은 라이브러리를 추가한다.

implementation 'io.awspring.cloud:spring-cloud-starter-aws:2.4.4'
implementation 'com.amazonaws:aws-java-sdk-sagemakerruntime:1.12.762'

 

 

2. SageMaker Runtime을 위한 Configuration을 등록한다.

package conference.clerker.global.aws;

import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.sagemakerruntime.AmazonSageMakerRuntime;
import com.amazonaws.services.sagemakerruntime.AmazonSageMakerRuntimeClientBuilder;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AmazonSageMakerRuntimeConfig {

    // AWS 자격 증명 및 리전 설정
    @Value("${other.aws.access-key}")
    private String accessKey;

    @Value("${other.aws.secret-key}")
    private String secretKey;

    @Value("${other.aws.region}")
    private String awsRegion;

    @Bean
    public AmazonSageMakerRuntime amazonSageMakerRuntime() {
        BasicAWSCredentials awsCreds = new BasicAWSCredentials(accessKey, secretKey);

        // 클라이언트 구성 설정
        ClientConfiguration clientConfig = new ClientConfiguration()
                .withConnectionTimeout(10_000) // 연결 타임아웃 설정 (예: 10초)
                .withSocketTimeout(300_000)    // 소켓 타임아웃 설정 (예: 5분)
                .withMaxErrorRetry(0);         // 자동 재시도 비활성화

        return AmazonSageMakerRuntimeClientBuilder.standard()
                .withCredentials(new AWSStaticCredentialsProvider(awsCreds))
                .withClientConfiguration(clientConfig)
                .withRegion(awsRegion)
                .build();
    }
}

 

 

3. SageMaker에 invoke 메서드를 활용하여 요청을 보내고 응답을 리턴한다.

 

@RestController
@RequiredArgsConstructor
public class ConnectToSageMakerEndpoint {

    private final AmazonSageMakerRuntime amazonSageMakerRuntime;

    @PostMapping("/invoke-sagemaker")
    public String invokeModelServerAsync(ModelRequestDTO modelRequestDTO) {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            String payload = objectMapper.writeValueAsString(modelRequestDTO);

            InvokeEndpointRequest invokeEndpointRequest = new InvokeEndpointRequest()
                    .withEndpointName(endpointName) // endpointName 직접 주입
                    .withContentType("application/json")
                    .withBody(ByteBuffer.wrap(payload.getBytes(StandardCharsets.UTF_8)));

            InvokeEndpointResult invokeEndpointResult = sagemakerRuntimeClient.invokeEndpoint(invokeEndpointRequest);
            
            return new String(invokeEndpointResult.getBody().array());
            
        } catch (Exception e) {
            log.error("모델 서버 호출 중 오류 발생: {}", e.getMessage(), e);
        }
    }
}

 

 

참고

https://docs.aws.amazon.com/ko_kr/sagemaker/latest/dg/serverless-endpoints-invoke.html

서버리스 엔드포인트 호출 - Amazon SageMaker

https://dummybot.medium.com/calling-sagemaker-runtime-endpoint-using-springboot-and-java-9ae11e11d2d7

Calling SageMaker Runtime Endpoint using SpringBoot and Java