Inconsistent behavior of argument model_dir between single instance vs. distributed training using TensorFlow Estimator

[shashankprasanna]

Describe the bug
The SageMaker Python SDK uses the argument model_dir inconsistently when executing single instance training vs. distributed training.

When running single instance training, SageMaker invokes specified python script with the argument:
python my_tf_script.py --model_dir s3://my_default_bucket/job/...

When running a distributed training job using horovod/MPI (by passing distributions to the TF estimator), sagemaker invokes MPIRUN with:
mpirun --host algo-1,algo-2 -np 2 /usr/local/bin/python3.6 -m mpi4py my_hvd_script.py --model_dir /opt/ml/model

This is very confusing and inconsistent behavior that's very hard to debug.

The SDK docs doesn't explain these differences:
https://sagemaker.readthedocs.io/en/stable/sagemaker.tensorflow.html
"model_dir (str) – S3 location where the checkpoint data and models can be exported to during training (default: None). If not specified a default S3 URI will be generated. It will be passed in the training script as one of the command line arguments."

The above definition is only true for single instance training. For Distributed training it's set to /opt/ml/model

To reproduce
To repro:
Run a single instance TensorFlow script mode training and observe logs. Here is an example where model_dir is set to an S3 location by the SDK.
https://github.com/shashankprasanna/reinvent-aim338/blob/master/2-dl-containers/cifar10_sagemaker_demo-with-results.ipynb
Output:
Invoking script with the following command:
/usr/local/bin/python3.6 cifar10-training-script-sagemaker.py --batch-size 256 --epochs 30 --learning-rate 0.01 --model_dir s3://reinvent-aim338/tensorflow-training-2019-12-05-19-13-24-229/model --momentum 0.9 --optimizer sgd --weight-decay 0.0002

Run distributed Tensorflow horovod/MPI script mode training and observe logs.
https://github.com/shashankprasanna/reinvent-aim338/blob/master/3-distributed-training/cifar10-sagemaker-distributed.ipynb
Output:
Command "mpirun --host algo-1,algo-2 -np 2 --allow-run-as-root <many_other_arguments> /usr/local/bin/python3.6 -m mpi4py cifar10-tf-horovod-sagemaker.py --batch-size 256 --epochs 100 --learning-rate 0.001 --model_dir /opt/ml/model --momentum 0.9 --optimizer adam --tensorboard_logs s3://sagemaker-jobs/jobstensorboard_logs/tf-horovod-2x1-workers-2020-03-13-00-25-10-073 --weight-decay 0.0002"

Expected behavior
When user doesn't override model_dir, SageMaker SDK must pass the same type of location (S3 or local) consistently for single instance or distributed training. Today one is S3 and the other is a local /opt/ml/model

Screenshots or logs
I've included logs/output in the repro section of this issue

System information
A description of your system. Please provide:

• SageMaker Python SDK version: '1.50.16'
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): TensorFlow
• Framework version: 1.14 (issue shows up on other versions 1.15, 1.13)
• Python version: python 3
• CPU or GPU: both
• Custom Docker image (Y/N): no, DLC images on ECR

Additional context
Add any other context about the problem here.

[laurenyu]

@shashankprasanna sorry for the delay in response, and sorry for the confusion

this is the expected behavior when using MPI for distributed training, so we'll need to update the documentation. however, if you use parameter servers, it should use an S3 location for model_dir. can you share your code for starting the training jobs?

[shashankprasanna]

Hi @laurenyu the training scripts are in the same repo as the Jupyter notebook code in the issue.

(PS: I deleted my earlier comment, since I first misread your comment as this is expected behavior and no changes were required.)

[laurenyu]

I've opened a PR to fix the documentation: #1368

[laurenyu]

released in v1.51.4