Skip to content

Give each host a unique rank in Hostfile.from_list - #4027

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
ayaangazali:fix-hostfile-rank-shadowing
Aug 6, 2026
Merged

Give each host a unique rank in Hostfile.from_list#4027
zcbenz merged 1 commit into
ml-explore:mainfrom
ayaangazali:fix-hostfile-rank-shadowing

Conversation

@ayaangazali

Copy link
Copy Markdown
Contributor

Proposed changes

Hostfile.from_list builds every Host with the wrong rank. The inner repeats loop reuses the name of the enumerate index:

for i, h in enumerate(hostlist.split(",")):
    ...
    for i in range(repeats):        # shadows the enumerate index
        hosts.append(Host(i, h, ips, []))

so the value passed to Host is the repeat counter, never the position in the list. With the default repeats=1 the range is always range(1), so every host gets rank 0.

from_file, the sibling constructor, assigns Host(i, ...) straight from enumerate, so a unique sequential rank is clearly the intent.

Repro

from mlx._distributed_utils.common import Hostfile

[(h.rank, h.ssh_hostname) for h in Hostfile.from_list("hostA,hostB,hostC").hosts]
# before: [(0, 'hostA'), (0, 'hostB'), (0, 'hostC')]
# after:  [(0, 'hostA'), (1, 'hostB'), (2, 'hostC')]

[(h.rank, h.ssh_hostname) for h in Hostfile.from_list("hostA,hostB", 3).hosts]
# before: [(0, 'hostA'), (1, 'hostA'), (2, 'hostA'), (0, 'hostB'), (1, 'hostB'), (2, 'hostB')]
# after:  [(0, 'hostA'), (1, 'hostA'), (2, 'hostA'), (3, 'hostB'), (4, 'hostB'), (5, 'hostB')]

Using the running length of the list keeps ranks unique across repeats too.

Scope

To be upfront: nothing reads Host.rank today, so this is latent rather than a live failure — I found it while reading the launcher, not from a bug report. It is still a wrong value on a public dataclass field that from_file populates correctly, so it seemed worth correcting before something starts depending on it. Close it if you would rather leave the field alone.

Behaviour that is unchanged: empty hostnames still raise, and literal IPs are still detected and stored in ips.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

The repeats loop reused the name of the enumerate index, so every Host
was constructed with the repeat counter instead of its position in the
list: with the default repeats=1 every host got rank 0. Track the rank
with the running length of the list, matching from_file.
@zcbenz
zcbenz merged commit d2275d5 into ml-explore:main Aug 6, 2026
28 checks passed
krystophny pushed a commit to krystophny/mlx-jaccl-fix-small-recv that referenced this pull request Aug 7, 2026
krystophny pushed a commit to krystophny/mlx-jaccl-fix-small-recv that referenced this pull request Aug 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants