Add mpiexecjl for windows

.github/workflows/ShellCheck.yml

@@ -26,4 +26,4 @@ jobs:
       - name: Install dependencies
         run: sudo apt install shellcheck
       - name: Check scripts
-        run: shellcheck bin/*
+        run: shellcheck bin/mpiexecjl_unix

bin/mpiexecjl_windows.ps1

@@ -0,0 +1,86 @@
+# Copyright (C) 2023 Guilherme Bodin
+# License is MIT "Expat"
+#
+# Commentary:
+#
+# Command line utility to call the `mpiexec` binary used by the `MPI.jl` version
+# in the given Julia project. It has the same syntax as the `mpiexec` binary
+# that would be called, with the additional `--project=...` flag to select a
+# different Julia project.
+#
+# Examples of usage (the MPI flags available depend on the MPI implementation
+# called):
+#
+#   $ mpiexecjl -n 40 julia mpi-script.jl
+#   $ mpiexecjl --project=my_experiment -n 80 --oversubscribe julia mpi-script.jl
+
+function usage {
+    Write-Host "Usage: $([System.IO.Path]::GetFileNameWithoutExtension($MyInvocation.MyCommand.Name)) [--project=...] MPIEXEC_ARGUMENTS..."
+    Write-Host "Call the mpiexec binary in the Julia environment specified by the --project option."
+    Write-Host "If no project is specified, the MPI associated with the global Julia environment will be used."
+    Write-Host "All other arguments are forwarded to mpiexec."
+}
+
+$julia_args = @()
+$PROJECT_ARG = ""
+
+echo $args
+
+foreach ($arg in $args) {
+    if ($arg -match "^--project(=.*)?$") {
+        $PROJECT_ARG = $arg
+    }
+    elseif ($arg -eq "-h" -or $arg -eq "--help") {
+        $helpRequested = $true
+        usage
+        Write-Host "Below is the help of the current mpiexec."
+        Write-Host
+        exit 0
+    }
+    else {
+        $julia_args += $arg
+    }
+}
+
+if (-not $julia_args) {
+    Write-Error "ERROR: no arguments specified."
+    usage
+    exit 1
+}
+
+if ($env:JULIA_BINDIR) {
+    $JULIA_CMD = Join-Path $env:JULIA_BINDIR "julia"
+} else {
+    $JULIA_CMD = "julia"
+}
+
+$SCRIPT = @'
+using MPI
+ENV[\"JULIA_PROJECT\"] = dirname(Base.active_project())
+try
+    mpiexec(exe -> run(`$exe $ARGS`))
+catch e
+    rethrow(e)
+end
+'@
+
+$PRECOMPILATION_SCRIPT = @'
+println(\"precompiling current project before running MPI\")
+
+import Pkg
+Pkg.activate(dirname(Base.active_project()))
+Pkg.instantiate()
+
+println(\"precompilation finished\")
+'@
+
+& $JULIA_CMD $PROJECT_ARG -e $PRECOMPILATION_SCRIPT
+
+if ($PROJECT_ARG) {
+    & $JULIA_CMD $PROJECT_ARG --color=yes --startup-file=no --compile=min -O0 -e $SCRIPT -- $julia_args
+} else {
+    & $JULIA_CMD --color=yes --startup-file=no --compile=min -O0 -e $SCRIPT -- $julia_args
+}
+
+# Guarantees that we will exit .ps1 with the same process status as the last command executed
+exit $lastExitCode

docs/src/usage.md

@@ -74,6 +74,13 @@ with:
 $ mpiexecjl --project=/path/to/project -n 20 julia script.jl
 ```
 
+On Windows systems, `mpiexecjl` is a powershell script that can be called as
+follows:
+
+```powershell
+powershell.exe -ExecutionPolicy ByPass -File $env:userprofile\\.julia\\bin\\mpiexecjl.ps1 --project=/path/to/project -n 20 julia script.jl
+```
+
 ## CUDA-aware MPI support
 
 If your MPI implementation has been compiled with CUDA support, then `CUDA.CuArray`s (from the

src/mpiexec_wrapper.jl

@@ -19,6 +19,13 @@ function install_mpiexecjl(; command::String = "mpiexecjl",
     end
     mkpath(destdir)
     verbose && @info "Installing `$(command)` to `$(destdir)`..."
-    cp(joinpath(@__DIR__, "..", "bin", "mpiexecjl"), exec; force = force)
+    if Sys.isunix()
+        cp(joinpath(@__DIR__, "..", "bin", "mpiexecjl_unix"), exec; force = force)
+    elseif Sys.iswindows()
+        exec *= ".ps1"
+        cp(joinpath(@__DIR__, "..", "bin", "mpiexecjl_windows.ps1"), exec; force = force)
+    else
+        throw(ErrorException("Unsupported platform: $(Sys.KERNEL)"))
+    end
     verbose && @info "Done!"
 end

test/mpiexecjl.jl

@@ -1,6 +1,88 @@
 using Test, Pkg
 using MPI
 
+function my_julia_cmd(julia=joinpath(Sys.BINDIR, Base.julia_exename()); cpu_target::Union{Nothing,String} = nothing)
+    opts = Base.JLOptions()
+    if cpu_target === nothing
+        cpu_target = unsafe_string(opts.cpu_target)
+    end
+    image_file = unsafe_string(opts.image_file)
+    addflags = String[]
+    let compile = if opts.compile_enabled == 0
+                      "no"
+                  elseif opts.compile_enabled == 2
+                      "all"
+                  elseif opts.compile_enabled == 3
+                      "min"
+                  else
+                      "" # default = "yes"
+                  end
+        isempty(compile) || push!(addflags, "--compile=$compile")
+    end
+    let depwarn = if opts.depwarn == 1
+                      "yes"
+                  elseif opts.depwarn == 2
+                      "error"
+                  else
+                      "" # default = "no"
+                  end
+        isempty(depwarn) || push!(addflags, "--depwarn=$depwarn")
+    end
+    let check_bounds = if opts.check_bounds == 1
+                      "yes" # on
+                  elseif opts.check_bounds == 2
+                      "no" # off
+                  else
+                      "" # default = "auto"
+                  end
+        isempty(check_bounds) || push!(addflags, "--check-bounds=$check_bounds")
+    end
+    opts.can_inline == 0 && push!(addflags, "--inline=no")
+    opts.use_compiled_modules == 0 && push!(addflags, "--compiled-modules=no")
+    opts.opt_level == 2 || push!(addflags, "-O$(opts.opt_level)")
+    opts.opt_level_min == 0 || push!(addflags, "--min-optlevel=$(opts.opt_level_min)")
+    push!(addflags, "-g$(opts.debug_level)")
+    if opts.code_coverage != 0
+        # Forward the code-coverage flag only if applicable (if the filename is pid-dependent)
+        coverage_file = (opts.output_code_coverage != C_NULL) ?  unsafe_string(opts.output_code_coverage) : ""
+        if isempty(coverage_file) || occursin("%p", coverage_file)
+            if opts.code_coverage == 1
+                push!(addflags, "--code-coverage=user")
+            elseif opts.code_coverage == 2
+                push!(addflags, "--code-coverage=all")
+            elseif opts.code_coverage == 3
+                push!(addflags, "--code-coverage=@$(unsafe_string(opts.tracked_path))")
+            end
+            isempty(coverage_file) || push!(addflags, "--code-coverage=$coverage_file")
+        end
+    end
+    if opts.malloc_log == 1
+        push!(addflags, "--track-allocation=user")
+    elseif opts.malloc_log == 2
+        push!(addflags, "--track-allocation=all")
+    elseif opts.malloc_log == 3
+        push!(addflags, "--track-allocation=@$(unsafe_string(opts.tracked_path))")
+    end
+    if opts.color == 1
+        push!(addflags, "--color=yes")
+    elseif opts.color == 2
+        push!(addflags, "--color=no")
+    end
+    if opts.startupfile == 2
+        push!(addflags, "--startup-file=no")
+    end
+    if opts.use_sysimage_native_code == 0
+        push!(addflags, "--sysimage-native-code=no")
+    end
+    if opts.use_pkgimages == 0
+        push!(addflags, "--pkgimages=no")
+    else
+        # If pkgimage is set, malloc_log and code_coverage should not
+        @assert opts.malloc_log == 0 && opts.code_coverage == 0
+    end
+    return `$julia -C$cpu_target --sysimage=$image_file $addflags`
+end
+
 @testset "mpiexecjl" begin
     mktempdir() do dir
         # Install MPI locally, so that we can test the `--project` flag to
@@ -15,24 +97,37 @@ using MPI
         # Test a run of mpiexec
         nprocs_str = get(ENV, "JULIA_MPI_TEST_NPROCS", "")
         nprocs = nprocs_str == "" ? clamp(Sys.CPU_THREADS, 2, 4) : parse(Int, nprocs_str)
-        mpiexecjl = joinpath(dir, "mpiexecjl")
+        mpiexecjl = if Sys.isunix()
+            joinpath(dir, "mpiexecjl")
+        elseif Sys.iswindows()
+            joinpath(dir, "mpiexecjl.ps1")
+        else
+            error("Unsupported platform: $(Sys.KERNEL)")
+        end
+
+        additional_initial_parts_cmd = if Sys.isunix()
+            String[]
+        else
+            String["powershell.exe", "-ExecutionPolicy", "Bypass", "-File"]
+        end
+
         # `Base.julia_cmd()` ensures keeping consistent flags when running subprocesses.
-        julia = Base.julia_cmd()
+        julia = my_julia_cmd()
         example = joinpath(@__DIR__, "..", "docs", "examples", "01-hello.jl")
         env = ["JULIA_BINDIR" => Sys.BINDIR]
         p = withenv(env...) do
-            run(`$(mpiexecjl) -n $(nprocs) --project=$(dir) $(julia) --startup-file=no -q $(example)`)
+            run(`$(additional_initial_parts_cmd) $(mpiexecjl) -n $(nprocs) --project=$(dir) $(julia) --startup-file=no -q $(example)`)
         end
         @test success(p)
         # Test help messages
         for help_flag in ("-h", "--help")
             help_message = withenv(env...) do
-                read(`$(mpiexecjl) --project=$(dir) --help`, String)
+                read(`$(additional_initial_parts_cmd) $(mpiexecjl) --project=$(dir) $help_flag`, String)
             end
             @test occursin(r"Usage:.*MPIEXEC_ARGUMENTS", help_message)
         end
         # Without arguments, or only with the `--project` option, the wrapper will fail
-        @test !withenv(() -> success(`$(mpiexecjl) --project=$(dir)`), env...)
-        @test !withenv(() -> success(`$(mpiexecjl)`), env...)
+        @test !withenv(() -> success(`$(additional_initial_parts_cmd) $(mpiexecjl) --project=$(dir)`), env...)
+        @test !withenv(() -> success(`$(additional_initial_parts_cmd) $(mpiexecjl)`), env...)
     end
 end

test/runtests.jl

@@ -60,11 +60,7 @@ if haskey(ENV,"JULIA_MPI_TEST_ABI")
     @test ENV["JULIA_MPI_TEST_ABI"] == MPIPreferences.abi
 end
 
-if Sys.isunix()
-    # This test doesn't need to be run with mpiexec.  `mpiexecjl` is currently
-    # available only on Unix systems
-    include("mpiexecjl.jl")
-end
+include("mpiexecjl.jl")
 
 excludefiles = split(get(ENV,"JULIA_MPI_TEST_EXCLUDE",""),',')