aws / aws/sagemaker-training-toolkit

Mpi mode sets all nodes to the same SM_CURRENT_HOST

Open
#158 0 comments 0 reactions 0 assignees View on GitHub
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.

Screenshot 2022-08-10 at 20 28 00

**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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.