feat: add OMPI cluster detection

Author: Avik Pal

ext/ReactantMPIExt.jl

@@ -3,9 +3,82 @@ module ReactantMPIExt
 using Reactant: Reactant, Distributed
 using MPI: MPI
 
-# Code taken from https://github.com/jax-ml/jax/blob/b0117366686ab084d38ad2657d9a2ae3a581ca7e/jax/_src/clusters/mpi4py_cluster.py
+# Code taken from:
+#   1. https://github.com/jax-ml/jax/blob/b0117366686ab084d38ad2657d9a2ae3a581ca7e/jax/_src/clusters/mpi4py_cluster.py
+#   2. https://github.com/jax-ml/jax/blob/b0117366686ab084d38ad2657d9a2ae3a581ca7e/jax/_src/clusters/ompi_cluster.py
 
-# XXX: Is this a good check??
+# Based on ompi_cluster
+const _ORTE_URI = "OMPI_MCA_orte_hnp_uri"
+const _PMIX_SERVER_URI = (
+    "PMIX_SERVER_URI2",
+    "PMIX_SERVER_URI3",
+    "PMIX_SERVER_URI4",
+    "PMIX_SERVER_URI41",
+    "PMIX_SERVER_URI21",
+)
+const _OMPI_PROCESS_COUNT = "OMPI_COMM_WORLD_SIZE"
+const _OMPI_PROCESS_ID = "OMPI_COMM_WORLD_RANK"
+const _OMPI_LOCAL_PROCESS_ID = "OMPI_COMM_WORLD_LOCAL_RANK"
+
+Distributed.is_env_present(::Distributed.OpenMPIORTEEnvDetector) = haskey(ENV, _ORTE_URI)
+
+function Distributed.is_env_present(::Distributed.OpenMPIPMIXEnvDetector)
+    return any(Base.Fix1(haskey, ENV), _PMIX_SERVER_URI)
+end
+
+function Distributed.get_coordinator_address(
+    ::Distributed.OpenMPIORTEEnvDetector, timeout_in_seconds::Integer
+)
+    orte_uri = ENV[_ORTE_URI]
+
+    job_id = parse(Int, split(orte_uri, '.'; limit=2)[1])
+    port = job_id % 2^12 + (65535 - 2^12 + 1)
+
+    launcher_ip_match = match(r"tcp://(.+?)[,:]|tcp6://\[(.+?)[,\]]", orte_uri)
+
+    @assert launcher_ip_match !== nothing "Could not parse coordinator IP address from \
+                                           Open MPI environment."
+
+    launcher_ip = launcher_ip_match.captures[findfirst(
+        !isnothing, launcher_ip_match.captures
+    )]
+    return "$(launcher_ip):$(port)"
+end
+
+function Distributed.get_coordinator_address(
+    ::Distributed.OpenMPIPMIXEnvDetector, timeout_in_seconds::Integer
+)
+    varname = findfirst(Base.Fix1(haskey, ENV), _PMIX_SERVER_URI)
+    pmix_uri = ENV[_PMIX_SERVER_URI[varname]]
+
+    job_id = parse(Int, split(split(pmix_uri, '-'; limit=3)[3], "@"; limit=2)[1])
+    port = job_id % 2^12 + (65535 - 2^12 + 1)
+
+    launcher_ip_match = match(r"tcp4://(.+?):|tcp6://\[(.+?)\]", pmix_uri)
+
+    @assert launcher_ip_match !== nothing "Could not parse coordinator IP address from \
+                                           Open MPI environment."
+
+    launcher_ip = launcher_ip_match.captures[findfirst(
+        !isnothing, launcher_ip_match.captures
+    )]
+
+    return "$(launcher_ip):$(port)"
+end
+
+function Distributed.get_process_count(::Distributed.AbstractOMPIClusterEnvDetector)
+    return parse(Int, ENV[_OMPI_PROCESS_COUNT])
+end
+
+function Distributed.get_process_id(::Distributed.AbstractOMPIClusterEnvDetector)
+    return parse(Int, ENV[_OMPI_PROCESS_ID])
+end
+
+function Distributed.get_local_process_id(::Distributed.AbstractOMPIClusterEnvDetector)
+    return parse(Int, ENV[_OMPI_LOCAL_PROCESS_ID])
+end
+
+# Based on mpi4py
 Distributed.is_env_present(::Distributed.MPIEnvDetector) = MPI.Initialized()
 
 function Distributed.get_coordinator_address(

src/Distributed.jl

@@ -36,6 +36,11 @@ end
 
 abstract type AbstractClusterEnvDetector end
 
+abstract type AbstractOMPIClusterEnvDetector <: AbstractClusterEnvDetector end
+
+struct OpenMPIORTEEnvDetector <: AbstractOMPIClusterEnvDetector end
+struct OpenMPIPMIXEnvDetector <: AbstractOMPIClusterEnvDetector end
+
 struct MPIEnvDetector <: AbstractClusterEnvDetector end
 
 # Based on https://github.com/jax-ml/jax/blob/b0117366686ab084d38ad2657d9a2ae3a581ca7e/jax/_src/clusters/cluster.py
@@ -48,7 +53,7 @@ function get_process_id end
 function get_local_process_id end
 
 function auto_detect_unset_distributed_params(;
-    detector_list=[MPIEnvDetector()],
+    detector_list=[OpenMPIORTEEnvDetector(), OpenMPIPMIXEnvDetector(), MPIEnvDetector()],
     coordinator_address::Union{Nothing,String}=nothing,
     num_processes::Union{Nothing,Integer}=nothing,
     process_id::Union{Nothing,Integer}=nothing,
@@ -70,6 +75,8 @@ function auto_detect_unset_distributed_params(;
 
     detector = detector_list[idx]
 
+    @debug "Detected cluster environment" detector
+
     if coordinator_address === nothing
         coordinator_address = get_coordinator_address(
             detector, initialization_timeout_in_seconds