Feature proposal: `Session` can download from S3

[andremoeller]

Please fill out the form below.

Describe the problem

sagemaker.session.Session has an upload_data method that allows users to upload a local file or directory to S3:

sagemaker-python-sdk/src/sagemaker/session.py

Lines 137 to 165 in 6f3cf42

	def upload_data(self, path, bucket=None, key_prefix="data", extra_args=None):
	"""Upload local file or directory to S3.
	If a single file is specified for upload, the resulting S3 object key is
	``{key_prefix}/{filename}`` (filename does not include the local path, if any specified).
	If a directory is specified for upload, the API uploads all content, recursively,
	preserving relative structure of subdirectories. The resulting object key names are:
	``{key_prefix}/{relative_subdirectory_path}/filename``.
	Args:
	path (str): Path (absolute or relative) of local file or directory to upload.
	bucket (str): Name of the S3 Bucket to upload to (default: None). If not specified, the
	default bucket of the ``Session`` is used (if default bucket does not exist, the
	``Session`` creates it).
	key_prefix (str): Optional S3 object key name prefix (default: 'data'). S3 uses the
	prefix to create a directory structure for the bucket content that it display in
	the S3 console.
	extra_args (dict): Optional extra arguments that may be passed to the upload operation.
	Similar to ExtraArgs parameter in S3 upload_file function. Please refer to the
	ExtraArgs parameter documentation here:
	https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html#the-extraargs-parameter
	Returns:
	str: The S3 URI of the uploaded file(s). If a file is specified in the path argument,
	the URI format is: ``s3://{bucket name}/{key_prefix}/{original_file_name}``.
	If a directory is specified in the path argument, the URI format is
	``s3://{bucket name}/{key_prefix}``.
	"""

But there's no corresponding way to download files from using Session into a local directory. Training jobs put model artifacts in S3, and transform jobs put batch transform output in S3, and any job (or even an Endpoint) may output to S3 during its execution but right now, users have to use boto3 instead of sagemaker_session

Proposal

Add sagemaker.session.Session.download_data with the following signature and behavior:

def download_data(self, s3_uri, local_path):
    """Downloads data under an S3 prefix from an S3 URI into a local path or directory.

    Args:
        s3_uri (str): An S3 path. All objects under this prefix will be downloaded into the
          current directory.
        local_path (str): A local path. "." means the current directory. Directories will be
          created as needed.
    """

Thoughts / feedback on this proposal are welcome. Thanks!

[ChoiByungWook]

@andremoeller,

This looks good. Can you provide a small snippet/example of using the new method in a typical customer workflow?

Thanks!

[andremoeller]

Sure:

Download a trained model. TF, for example, contains a TF SavedModel, which can be inspected for its signature.

estimator.fit(..., wait=True)
# Downloads model.tar.gz into 'my/local/model_directory'
estimator.sagemaker_session.download_data(estimator.output_path, 'my/local/model_directory')

Download input and output from a transformed job and load it into a pandas DataFrame. This makes it easy to visualize model performance, do model evaluation, see how features affect, and combine input record indices with transform output (in case the model only gives predictions, and not record indices).

def get_dataframe_from_csv_files(dir):
    csv_files = os.listdir(transform_out)
    dfs = (pd.read_csv(f) for f in all_files)
    return pd.concat(df_from_each_file, ignore_index=True)

transformer.transform(...)
transformer.wait()
# Downloads output from transformer's output path into 'transform_out'
transformer.sagemaker_session.download_data(transformer.output_path, 'transform_out')

# Read output into a pandas dataframe:
output df = get_dataframe_from_csv_files('transform_out')

# Download and read input into a pandas dataframe:
transformer.sagemaker_session.download_data('s3://input-bucket/csv/data', 'transform_in')
input_df = pd.read_csv('transform_in')

# Create one pandas DataFrame with input and output:
combined_df = pd.concat([input_df, output_df], axis=1) 

This would make the 'boto3' section of SageMaker notebooks like this one, where data is downloaded using boto3 rather than the SageMaker Python SDK, unnecessary:

https://github.com/awslabs/amazon-sagemaker-examples/blob/81167a0d1fad2f433766117980513413c447f6b5/sagemaker-python-sdk/tensorflow_distributed_mnist/tensorflow_batch_transform_mnist.ipynb

Raw link in case the link above can't be rendered:
https://raw.githubusercontent.com/awslabs/amazon-sagemaker-examples/80333fd4632cf6d924d0b91c33bf80da3bdcf926/introduction_to_amazon_algorithms/imageclassification_caltech/Image-classification-fulltraining.ipynb#L110-115

[laurenyu]

there is now an S3Downloader object that provides this functionality.

docs: https://sagemaker.readthedocs.io/en/stable/s3.html#sagemaker.s3.S3Downloader