Use PyTorch Estimator for SageMaker training jobs

[AdeelH]

Overview

This PR replaces the use of the generic Estimator with the PyTorch-specific Estimator provided by the SageMaker Python SDK. This allows us to distribute training jobs across multiple instances using torchrun.

Other changes:

• Add train_image, train_instance_type, and train_instance_count RV config options.

Checklist

• Added unit tests, if applicable
• Updated documentation, if applicable
• Added needs-backport label if the change should be back-ported to the previous release
• PR has a name that won't get you publicly shamed for vagueness

Notes

The current hack of creating a fake python script will not be necessary once aws/sagemaker-python-sdk#4324 is merged and released.

Testing Instructions

Sample command:

SAGEMAKER_TRAIN_INSTANCE_TYPE=ml.p3.8xlarge SAGEMAKER_TRAIN_INSTANCE_COUNT=2 rastervision run sagemaker \
	"rastervision_pytorch_backend/rastervision/pytorch_backend/examples/semantic_segmentation/isprs_potsdam.py" \
	-a raw_uri "s3://raster-vision-raw-data/isprs-potsdam" \
	-a root_uri "s3://raster-vision-ahassan/rvtest/sagemaker_dist/isprs_potsdam_2024_01_23b/"

[codecov]

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (6a918ba) 85.16% compared to head (1752019) 85.23%.

Files	Patch %	Lines
...rastervision/aws_sagemaker/aws_sagemaker_runner.py	98.48%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2031      +/-   ##
==========================================
+ Coverage   85.16%   85.23%   +0.06%     
==========================================
  Files         196      196              
  Lines        9856     9908      +52     
==========================================
+ Hits         8394     8445      +51     
- Misses       1462     1463       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.