-
Notifications
You must be signed in to change notification settings - Fork 2.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
commit 79b8cbf Author: Corey Adams <[email protected]> Date: Mon Jul 1 14:14:15 2024 -0500 Fix mypy issues; change variable name to more universally known name commit 10edc86 Author: Corey Adams <[email protected]> Date: Thu Jun 27 13:25:32 2024 -0500 Change copyright year to the year this was authored commit f7086cb Author: Corey Adams <[email protected]> Date: Thu Jun 27 13:15:32 2024 -0500 Update build file to include mpi4py cluster. commit 6235eb3 Author: Corey adams <[email protected]> Date: Thu Jun 27 12:11:48 2024 -0500 Update distributed.py Clean up documentation slightly. commit ef3a2e2 Author: Corey adams <[email protected]> Date: Thu Jun 27 12:09:37 2024 -0500 Update mpi4py_cluster.py Further clean up unneeded comments. commit 6cc07a9 Author: Corey adams <[email protected]> Date: Thu Jun 27 12:08:38 2024 -0500 Update mpi4py_cluster.py Remove unneeded commented code. commit 6701bd1 Merge: 5a91ac3 98b8754 Author: Corey adams <[email protected]> Date: Thu Jun 27 12:07:25 2024 -0500 Merge branch 'google:main' into main commit 5a91ac3 Merge: 301bbc6 6c51234 Author: Corey adams <[email protected]> Date: Tue May 28 22:14:08 2024 -0500 Merge branch 'google:main' into main commit 301bbc6 Author: Corey Adams <[email protected]> Date: Tue May 28 11:34:51 2024 -0500 Add test to verify mpi4py based distributed initialization commit 19e6694 Author: Corey Adams <[email protected]> Date: Tue May 28 11:14:40 2024 -0500 Unify variable naming and fix function argument ordering commit 72fe093 Author: Corey Adams <[email protected]> Date: Tue May 28 10:56:25 2024 -0500 Remove unmerged code commit 3a96e73 Merge: e4fd97e ff3db9b Author: Corey adams <[email protected]> Date: Tue May 28 10:51:41 2024 -0500 Merge branch 'google:main' into main commit e4fd97e Merge: a697299 72a81e5 Author: Corey adams <[email protected]> Date: Mon May 13 16:01:35 2024 -0500 Merge branch 'google:main' into main commit a697299 Merge: 85bcf42 1e48adc Author: Corey adams <[email protected]> Date: Mon May 13 14:21:32 2024 -0500 Merge branch 'google:main' into main commit 85bcf42 Merge: af1a4f0 06cd05d Author: Corey Adams <[email protected]> Date: Tue Apr 16 09:09:31 2024 -0500 Merge branch 'main' of https://github.com/google/jax commit af1a4f0 Author: Corey Adams <[email protected]> Date: Tue Apr 16 08:58:33 2024 -0500 update documentation and elaborate on spec_detect_method variable commit 01f4709 Author: Corey Adams <[email protected]> Date: Tue Apr 16 08:45:38 2024 -0500 Address feedback and comments on PR 20174; fix typo in documentation. commit 4f22d86 Merge: 900a037 71ec6e3 Author: Corey adams <[email protected]> Date: Mon Mar 11 11:51:30 2024 -0500 Merge branch 'google:main' into main commit 900a037 Author: Corey Adams <[email protected]> Date: Mon Mar 11 11:50:48 2024 -0500 Auto-detect of mpi4py-based configuration is now strictly opt-in. commit 1992969 Author: Corey Adams <[email protected]> Date: Thu Mar 7 12:27:43 2024 -0600 Enable automatic detection of distrbuted variables with any configuration of MPI, as long as mpi4py is available
- Loading branch information
Showing
9 changed files
with
208 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
# Copyright 2024 The JAX Authors. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
# | ||
# https://www.apache.org/licenses/LICENSE-2.0 | ||
# | ||
# Unless required by applicable law or agreed to in writing, software | ||
# distributed under the License is distributed on an "AS IS" BASIS, | ||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
# See the License for the specific language governing permissions and | ||
# limitations under the License. | ||
|
||
from __future__ import annotations | ||
|
||
from jax._src import clusters | ||
import socket | ||
|
||
from importlib.util import find_spec | ||
|
||
|
||
class Mpi4pyCluster(clusters.ClusterEnv): | ||
|
||
|
||
name: str = "mpi4py" | ||
opt_in_only_method: bool = True | ||
|
||
@classmethod | ||
def is_env_present(cls) -> bool: | ||
|
||
# Relies on mpi4py: | ||
return find_spec("mpi4py") is not None | ||
|
||
@classmethod | ||
def get_coordinator_address(cls, timeout_secs: int | None) -> str: | ||
|
||
# Using mpi4py, figure out rank 0 and it's hostname. | ||
# Then broadcast the hostname and port. | ||
|
||
|
||
from mpi4py import MPI #type: ignore | ||
# Get the global communicator: | ||
COMM_WORLD = MPI.COMM_WORLD | ||
|
||
# On rank 0, get the hostname: | ||
|
||
if COMM_WORLD.Get_rank() == 0: | ||
# Order all the hostnames, and find unique ones | ||
hostname = socket.gethostname() | ||
|
||
# Apparently, we want to pick a port in an ephemeral range... | ||
port_id = hash(hostname) % 2**12 + (65535 - 2**12 + 1) | ||
|
||
hostname = f'{hostname}:{port_id}' | ||
|
||
else: | ||
hostname = "None" | ||
|
||
|
||
|
||
# Broadcast the host_ip to all ranks: | ||
hostname = COMM_WORLD.bcast(hostname, root=0) | ||
|
||
|
||
return hostname | ||
|
||
|
||
@classmethod | ||
def get_process_count(cls) -> int: | ||
from mpi4py import MPI | ||
return int(MPI.COMM_WORLD.Get_size()) | ||
|
||
@classmethod | ||
def get_process_id(cls) -> int: | ||
from mpi4py import MPI | ||
return int(MPI.COMM_WORLD.Get_rank()) | ||
|
||
@classmethod | ||
def get_local_process_id(cls) -> int | None: | ||
|
||
# Using mpi4py, split the global communicator into sub communicators | ||
# based on hostname. mpi will assign them ranks and that will allow | ||
# a selection of the local process ID. | ||
from mpi4py import MPI | ||
COMM_WORLD = MPI.COMM_WORLD | ||
|
||
# This is the alternative method that is simpler: | ||
new_comm = COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED) | ||
|
||
|
||
# The rank in the new communicator - which is host-local only - IS the local rank: | ||
return int(new_comm.Get_rank()) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters