aws / aws/sagemaker-python-sdk

Add ml.p5e.48xlarge to EFA instance lists in sagemaker-train and sagemaker-core

Open
#5,491 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.3k
Forks
1.3k
Avg merge
1d 22h
Merged PRs (30d)
35

Description

# Add ml.p5e.48xlarge to EFA instance lists in sagemaker-train and sagemaker-core

## Description

The `SM_EFA_NCCL_INSTANCES` and `SM_EFA_RDMA_INSTANCES` lists in the sagemaker-python-sdk are missing `ml.p5e.48xlarge`, causing NCCL hangs during distributed training initialization on P5e instances when using the SDK's container drivers.

Additionally, `ml.p5.48xlarge` is missing from `SM_EFA_RDMA_INSTANCES` (it's only in `SM_EFA_NCCL_INSTANCES`).

## Current State

```python
SM_EFA_NCCL_INSTANCES = [
"ml.g4dn.8xlarge",
"ml.g4dn.12xlarge",
"ml.g5.48xlarge",
"ml.p3dn.24xlarge",
"ml.p4d.24xlarge",
"ml.p4de.24xlarge",
"ml.p5.48xlarge",
"ml.trn1.32xlarge",
]

SM_EFA_RDMA_INSTANCES = [
"ml.p4d.24xlarge",
"ml.p4de.24xlarge",
"ml.trn1.32xlarge",
]
```

## Expected State

```python
SM_EFA_NCCL_INSTANCES = [
"ml.g4dn.8xlarge",
"ml.g4dn.12xlarge",
"ml.g5.48xlarge",
"ml.p3dn.24xlarge",
"ml.p4d.24xlarge",
"ml.p4de.24xlarge",
"ml.p5.48xlarge",
"ml.p5e.48xlarge", # ADD
"ml.trn1.32xlarge",
]

SM_EFA_RDMA_INSTANCES = [
"ml.p4d.24xlarge",
"ml.p4de.24xlarge",
"ml.p5.48xlarge", # ADD
"ml.p5e.48xlarge", # ADD
"ml.trn1.32xlarge",
]
```

## Impact

Without these entries, the SDK's container drivers don't set the required EFA environment variables (`FI_PROVIDER=efa`, `FI_EFA_USE_DEVICE_RDMA=1`, `RDMAV_FORK_SAFE=1`) for P5e instances, causing NCCL to hang during collective initialization in multi-node distributed training.

## Related

- sagemaker-training-toolkit issue: https://github.com/aws/sagemaker-training-toolkit/issues/240
- sagemaker-training-toolkit PR: https://github.com/aws/sagemaker-training-toolkit/pull/241
- P5e instances use EFA with RDMA support, same as P4d/P4de/P5

## Questions

1. Is there a specific process for testing EFA/instance-specific changes on actual hardware before merging?
2. Should integration tests be added for P5e EFA configuration, or are unit tests sufficient?

Contributor guide

Open the contributing guide

Research direction

Start by locating the SM_EFA_NCCL_INSTANCES and SM_EFA_RDMA_INSTANCES definitions in the sagemaker-train and sagemaker-core components. Compare both lists with the expected state in this issue, then check the relevant unit-test coverage; done means both P5 instances are represented in the appropriate EFA lists.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python
Domain
cloud, machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.