test: test SimpleChains and Flux as separate projects (#1474)

* test: test SimpleChains and Flux as separate projects

* test: fixes

* test: move distributed testing deps into a project

* chore: run fmt

* test: missing test

* ci: separate out Flux/SimpleChains and Distributed test group

* test: fix tests

Author: avik-pal

.buildkite/testing.yml

@@ -37,6 +37,7 @@ steps:
               - "autodiff"
               - "misc"
               - "reactant"
+              - "extras"
 
   # - group: ":julia: (Lux) AMD GPU"
   #   steps:

.github/workflows/CI.yml

@@ -40,6 +40,7 @@ jobs:
           - "autodiff"
           - "misc"
           - "reactant"
+          - "extras"
     uses: ./.github/workflows/CommonCI.yml
     with:
       julia_version: ${{ matrix.version }}
@@ -60,6 +61,7 @@ jobs:
           - "autodiff"
           - "misc"
           - "reactant"
+          - "extras"
     uses: ./.github/workflows/CommonCI.yml
     with:
       julia_version: "1.11"

src/transform/simplechains.jl

@@ -25,7 +25,7 @@ performed using SimpleChains.
 
 ## Example
 
-```jldoctest
+```julia
 julia> import SimpleChains
 
 julia> using Adapt, Lux, Random

test/Project.toml

@@ -19,7 +19,6 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
-MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
@@ -34,7 +33,6 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
-SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -70,7 +68,6 @@ LuxCore = "1.2"
 LuxLib = "1.3.4"
 LuxTestUtils = "2"
 MLDataDevices = "1.6.10"
-MLUtils = "0.4.3"
 Mooncake = "0.4.138"
 NNlib = "0.9.26"
 Octavian = "0.3.28"
@@ -85,7 +82,6 @@ Reactant = "0.2.160"
 Reexport = "1.2.2"
 ReverseDiff = "1.15.3"
 Setfield = "1.1.1"
-SimpleChains = "0.4.7"
 StableRNGs = "1.0.2"
 Static = "1"
 StaticArrays = "1.9"

test/distributed/Project.toml

@@ -0,0 +1,29 @@
+[deps]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+NCCL = "3fe64909-d7a1-4096-9b7d-7a0f12cf0f6b"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Lux = {path = "../.."}
+LuxCUDA = {path = "../../lib/LuxCUDA"}
+MLDataDevices = {path = "../../lib/MLDataDevices"}
+
+[compat]
+ComponentArrays = "0.15.29"
+Lux = "1.21"
+LuxCUDA = "0.3.3"
+MLDataDevices = "1.12"
+MLUtils = "0.4.8"
+MPI = "0.20.23"
+NCCL = "0.1.1"
+Optimisers = "0.4.6"
+Pkg = "1.10"
+Test = "1.10"
+

test/distributed/common_distributedtest.jl

@@ -1,14 +1,20 @@
 using Lux, MPI, Test
 
 const input_args = length(ARGS) == 2 ? ARGS : ("cpu", "mpi")
+
 if input_args[1] == "cuda"
-    using LuxCUDA, NCCL
+    using LuxCUDA
 end
 if input_args[1] == "amdgpu"
     using AMDGPU
 end
 
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+
+if input_args[1] == "nccl"
+    using NCCL
+end
+
 const dev = if input_args[1] == "cpu"
     CPUDevice()
 else

test/distributed/data_distributedtest.jl

@@ -7,14 +7,20 @@ using Lux, MPI, Random, Test
 using MLUtils
 
 const input_args = length(ARGS) == 2 ? ARGS : ("cpu", "mpi")
+
 if input_args[1] == "cuda"
-    using LuxCUDA, NCCL
+    using LuxCUDA
 end
 if input_args[1] == "amdgpu"
     using AMDGPU
 end
 
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+
+if input_args[1] == "nccl"
+    using NCCL
+end
+
 const dev = if input_args[1] == "cpu"
     CPUDevice()
 else

test/distributed/distributed_test_runner.jl

@@ -0,0 +1,40 @@
+using MPI, Pkg, Test
+
+nprocs_str = get(ENV, "JULIA_MPI_TEST_NPROCS", "")
+nprocs = nprocs_str == "" ? clamp(Sys.CPU_THREADS, 2, 4) : parse(Int, nprocs_str)
+
+testdir = @__DIR__
+isdistributedtest(f) = endswith(f, "_distributedtest.jl")
+
+distributedtestfiles = String[]
+for (root, dirs, files) in walkdir(testdir)
+    for file in files
+        if isdistributedtest(file)
+            push!(distributedtestfiles, joinpath(root, file))
+        end
+    end
+end
+
+cur_proj = dirname(Pkg.project().path)
+
+@info "Running Distributed Tests with $nprocs processes"
+
+include("../setup_modes.jl")
+
+@testset "distributed tests: $(mode)" for (mode, aType, dev, ongpu) in MODES
+    backends = mode == "cuda" ? ("mpi", "nccl") : ("mpi",)
+    @testset "Backend: $(backend_type)" for backend_type in backends
+        np = backend_type == "nccl" ? min(nprocs, length(CUDA.devices())) : nprocs
+        @testset "$(basename(file))" for file in distributedtestfiles
+            @info "Running $file with $backend_type backend on $mode device"
+            try
+                run(`$(MPI.mpiexec()) -n $(np) $(Base.julia_cmd()) --color=yes \
+                     --code-coverage=user --project=$(cur_proj) --startup-file=no \
+                     $(file) $(mode) $(backend_type)`)
+                @test true
+            catch
+                @test false
+            end
+        end
+    end
+end

test/distributed/optimizer_distributedtest.jl

@@ -1,14 +1,20 @@
 using Lux, MPI, Optimisers, Random, Test
 
 const input_args = length(ARGS) == 2 ? ARGS : ("cpu", "mpi")
+
 if input_args[1] == "cuda"
-    using LuxCUDA, NCCL
+    using LuxCUDA
 end
 if input_args[1] == "amdgpu"
     using AMDGPU
 end
 
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+
+if input_args[1] == "nccl"
+    using NCCL
+end
+
 const dev = if input_args[1] == "cpu"
     CPUDevice()
 else

test/distributed/synchronize_distributedtest.jl

@@ -1,14 +1,20 @@
 using ComponentArrays, Lux, MPI, Optimisers, Random, Test
 
 const input_args = length(ARGS) == 2 ? ARGS : ("cpu", "mpi")
+
 if input_args[1] == "cuda"
-    using LuxCUDA, NCCL
+    using LuxCUDA
 end
 if input_args[1] == "amdgpu"
     using AMDGPU
 end
 
 const backend_type = input_args[2] == "nccl" ? NCCLBackend : MPIBackend
+
+if input_args[1] == "nccl"
+    using NCCL
+end
+
 const dev = if input_args[1] == "cpu"
     CPUDevice()
 else