Fix GreedyScheduler with ntasks and chunksize, closes #120 (#122)

JuliaFolds2 · Sep 25, 2024 · af70b5d · af70b5d
1 parent 8ea0ff7
commit af70b5d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 18 deletions.
diff --git a/src/schedulers.jl b/src/schedulers.jl
@@ -281,12 +281,13 @@ function GreedyScheduler(;
     else
         # only choose nchunks default if chunksize hasn't been specified
         if !isgiven(nchunks) && !isgiven(chunksize)
+            # neither nchunks nor chunksize has been specified but chunking = true
             nchunks = 10 * nthreads(:default)
             chunksize = -1
-        else
-            nchunks = isgiven(nchunks) ? nchunks :
-                      isgiven(ntasks) ? ntasks : -1
-            chunksize = isgiven(chunksize) ? chunksize : -1
+        elseif !isgiven(nchunks)
+            nchunks = -1
+        elseif !isgiven(chunksize)
+            chunksize = -1
         end
     end
     GreedyScheduler(ntasks, nchunks, chunksize, split; chunking)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -318,7 +318,7 @@ end;
 end;
 
 @testset "chunking mode + chunksize option" begin
-    for sched in (DynamicScheduler, StaticScheduler)
+    for sched in (DynamicScheduler, StaticScheduler, GreedyScheduler)
         @test sched() isa sched
         @test sched(; chunksize = 2) isa sched
 
@@ -328,26 +328,41 @@ end;
               OhMyThreads.Schedulers.FixedCount
         @test OhMyThreads.Schedulers.chunking_mode(sched(; chunking = false)) ==
               OhMyThreads.Schedulers.NoChunking
-        @test OhMyThreads.Schedulers.chunking_mode(sched(;
-            nchunks = 2, chunksize = 4, chunking = false)) ==
-              OhMyThreads.Schedulers.NoChunking
-        @test OhMyThreads.Schedulers.chunking_mode(sched(;
-            nchunks = -2, chunksize = -4, split = :whatever, chunking = false)) ==
-              OhMyThreads.Schedulers.NoChunking
+        if sched != GreedyScheduler
+            # For (Dynamic|Static)Scheduler `chunking = false` overrides `nchunks` and
+            # `chunksize`
+            @test OhMyThreads.Schedulers.chunking_mode(sched(;
+                nchunks = 2, chunksize = 4, chunking = false)) ==
+                  OhMyThreads.Schedulers.NoChunking
+            @test OhMyThreads.Schedulers.chunking_mode(sched(;
+                nchunks = -2, chunksize = -4, split = :whatever, chunking = false)) ==
+                  OhMyThreads.Schedulers.NoChunking
+            @test OhMyThreads.Schedulers.chunking_enabled(sched(;
+                nchunks = -2, chunksize = -4, chunking = false)) == false
+            @test OhMyThreads.Schedulers.chunking_enabled(sched(;
+                nchunks = 2, chunksize = 4, chunking = false)) == false
+        else
+            # For GreedyScheduler `nchunks` or `chunksize` overrides `chunking = false`
+            @test OhMyThreads.Schedulers.chunking_mode(sched(;
+                nchunks = 2, chunking = false)) ==
+                  OhMyThreads.Schedulers.FixedCount
+            @test OhMyThreads.Schedulers.chunking_mode(sched(;
+                chunksize = 2, chunking = false)) ==
+                  OhMyThreads.Schedulers.FixedSize
+            @test OhMyThreads.Schedulers.chunking_enabled(sched(;
+                nchunks = 2, chunking = false)) == true
+            @test OhMyThreads.Schedulers.chunking_enabled(sched(;
+                chunksize = 4, chunking = false)) == true
+        end
         @test OhMyThreads.Schedulers.chunking_enabled(sched(; chunksize = 2)) == true
         @test OhMyThreads.Schedulers.chunking_enabled(sched(; nchunks = 2)) == true
-        @test OhMyThreads.Schedulers.chunking_enabled(sched(;
-            nchunks = -2, chunksize = -4, chunking = false)) == false
-        @test OhMyThreads.Schedulers.chunking_enabled(sched(;
-            nchunks = 2, chunksize = 4, chunking = false)) == false
-
         @test_throws ArgumentError sched(; nchunks = 2, chunksize = 3)
         @test_throws ArgumentError sched(; nchunks = 0, chunksize = 0)
         @test_throws ArgumentError sched(; nchunks = -2, chunksize = -3)
 
-        let scheduler = sched(; chunksize = 2)
+        let scheduler = sched(; chunksize = 2, split = :batch)
             @test tmapreduce(sin, +, 1:10; scheduler) ≈ mapreduce(sin, +, 1:10)
-            @test tmap(sin, 1:10; scheduler) ≈ map(sin, 1:10)
+            @test tmap(sin, Float64, 1:10; scheduler) ≈ map(sin, 1:10)
             @test isnothing(tforeach(sin, 1:10; scheduler))
             @test treduce(+, 1:10; scheduler) ≈ reduce(+, 1:10)
         end