FSx for lustre with AWS Sagemaker Error - Artifact upload failed:Please ensure that the subnet's route table has a route to an S3 VPC endpoint or a NAT device

Hi,

I am trying to use FSx for Lustre for my Sagemaker training. I followed this tutorial: https://aws.amazon.com/blogs/machine-learning/speed-up-training-on-amazon-sagemaker-using-amazon-efs-or-amazon-fsx-for-lustre-file-systems/.

Code:

import sagemaker
from sagemaker import get_execution_role
from sagemaker.estimator import Estimator

sess = sagemaker.Session()

estimator = Estimator(image_uri='image-uri',
                      role='my-role',
                      base_job_name='training-job',
                      instance_count=1,
                      sagemaker_session=sess,
                      instance_type='ml.m5.xlarge',  # ml.p3.2xlarge
                      subnets=['subnet-id'],
                      security_group_ids=['sg-id', 'sg-id'],
                      hyperparameters={...} 
                     )

from sagemaker.inputs import FileSystemInput

train_input = FileSystemInput(file_system_id = 'fs-id',
                              file_system_type = 'FSxLustre',
                              directory_path = '/mount-idea/dataset',
                              file_system_access_mode = 'rw')

estimator.fit(train_input)

First issue is that I don't have internet connection in the container, and second is this error:

UnexpectedStatusException: Error for Training job training-job: Failed. Reason: ClientError: Artifact upload failed:Please ensure that the subnet's route table has a route to an S3 VPC endpoint or a NAT device, and both the security groups and the subnet's network ACL allow uploading data to all output URIs

I have create a custom VPC for this with private and public subnets. I enabled NAT and S3 endpoint. I am using one of the public subnet so that I can have access to internet. I created custom security group rules (inbound/outbound) with Custom TCP for port 988 and ports 1018 - 1023 for the security groups I use.

The code is running in a sagemaker notebook that has as VPC, the custom VPC I created and uses same subnet and security groups as the ones pass to the estimator.

What should I do to fix this error?

Thank you!