aws / aws/sagemaker-training-toolkit
Mpi mode sets all nodes to the same SM_CURRENT_HOST
- Dominant language
- Python
- Stars
- 530
- Forks
- 140
- Avg merge
- 1h 12m
- Merged PRs (30d)
- 2
Description
**Describe the bug**
With mpi mode, all nodes report the same `SM_CURRENT_HOST` (which is the master's one).
**To reproduce**
Run an PyTorch estimator in mpi mode and more than one node. The training entrypoint can simply dump all its environment variables to stdout (which should end-up on Cloudwatch log). From there, we can see that `SM_CURRENT_HOST` from all nodes are set to the same value (i.e., the master's), whereas `PMIX_HOSTNAME` is set correctly.
**Expected behavior**
Master node should not propagate its `SM_CURRENT_HOST` to the other nodes.
**Screenshots or logs**
If applicable, add screenshots or logs to help explain your problem.

**System information**
PyTorch DLC 1.11.0-gpu-py38
**Additional context**
Add any other context about the problem here.
This patch corrected the `SM_CURRENT_HOST` issue on my training jobs.
```python
# https://github.com/aws/sagemaker-training-toolkit/blob/3188a9df7803798defb043a332d789f7474219d0/src/sagemaker_training/mpi.py#L353
for name in self._env_vars:
if name.startswith("SM_"): # New addition
continue # New addition
command.extend(["-x", name])
```
Contributor guide
Research direction
Start at src/sagemaker_training/mpi.py around line 353 and inspect how environment variables are exported for MPI workers. Reproduce with a multi-node PyTorch estimator and compare SM_CURRENT_HOST with PMIX_HOSTNAME; done when each node reports its own host without breaking other environment propagation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100