Default partition in Distributed (#3496)

* correct partition * inothing -> isnothing * does not work on julia 1.9? * hopefully last one? * this was the last * other ones slipped * other ones remaining --------- Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Mar 25, 2024 · fb9455b · fb9455b
1 parent 7c2c0e7
commit fb9455b
Show file tree

Hide file tree

Showing 10 changed files with 10 additions and 10 deletions.
diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl
@@ -1,6 +1,6 @@
 using Oceananigans.Fields: OneField
 using Oceananigans.Grids: architecture
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 import Oceananigans.Fields: condition_operand, conditional_length, set!, compute_at!, indices
 
 # For conditional reductions such as mean(u * v, condition = u .> 0))

diff --git a/src/DistributedComputations/distributed_architectures.jl b/src/DistributedComputations/distributed_architectures.jl
@@ -204,16 +204,19 @@ Keyword arguments
                   if not for testing or developing. Change at your own risk!
 """
 function Distributed(child_architecture = CPU(); 
-                     communicator = MPI.COMM_WORLD,
+                     communicator = nothing,
                      devices = nothing, 
                      synchronized_communication = false,
-                     partition = Partition(MPI.Comm_size(communicator)))
+                     partition = nothing)
 
     if !(MPI.Initialized())
         @info "MPI has not been initialized, so we are calling MPI.Init()."
         MPI.Init()
     end
 
+    communicator = isnothing(communicator) ? MPI.COMM_WORLD : communicator
+    partition    = isnothing(partition) ? Partition(MPI.Comm_size(communicator)) : partition
+
     ranks = size(partition)
     Rx, Ry, Rz = ranks
     total_ranks = Rx * Ry * Rz

diff --git a/src/DistributedComputations/distributed_on_architecture.jl b/src/DistributedComputations/distributed_on_architecture.jl
@@ -1,5 +1,6 @@
 using CUDA: CuArray
 using OffsetArrays
+import Oceananigans.Architectures: on_architecture
 
 # We do not support switching from distributed and serial through `on_architecture`.
 # We only support moving a type from CPU to GPU and the other way around

diff --git a/src/DistributedComputations/partition_assemble.jl b/src/DistributedComputations/partition_assemble.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 
 all_reduce(op, val, arch::Distributed) = 
     MPI.Allreduce(val, op, arch.communicator)

diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl
@@ -4,7 +4,6 @@ using OffsetArrays: OffsetArray
 using Oceananigans.Utils: getnamewrapper
 using Oceananigans.Grids: total_size
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Oceananigans.BoundaryConditions: FBC
 using Printf
 

diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl
@@ -1,6 +1,5 @@
 using KernelAbstractions: @kernel, @index
 using Statistics
-using Oceananigans.Architectures: architecture
 using Oceananigans.Fields: location, ZReducedField, Field
 
 instantiate(T::Type) = T()

diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl
@@ -1,6 +1,5 @@
 using Oceananigans.Utils: prettysummary
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Printf
 
 #####

diff --git a/src/TurbulenceClosures/implicit_explicit_time_discretization.jl b/src/TurbulenceClosures/implicit_explicit_time_discretization.jl
@@ -1,4 +1,3 @@
-using Oceananigans.Utils: on_architecture
 using Oceananigans.Grids: AbstractGrid
 
 abstract type AbstractTimeDiscretization end

diff --git a/...Closures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl b/...Closures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators: ℑzᵃᵃᶜ

diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators
 using Oceananigans.Grids: inactive_node