Add support for limiting the evaluations of a streaming DAG

This is useful for testing and benchmarking.

Author: JamesWrigley

src/sch/Sch.jl

@@ -253,9 +253,11 @@ end
 Combine `SchedulerOptions` and `ThunkOptions` into a new `ThunkOptions`.
 """
 function Base.merge(sopts::SchedulerOptions, topts::ThunkOptions)
-    single = topts.single !== nothing ? topts.single : sopts.single
-    allow_errors = topts.allow_errors !== nothing ? topts.allow_errors : sopts.allow_errors
-    proclist = topts.proclist !== nothing ? topts.proclist : sopts.proclist
+    select_option = (sopt, topt) -> isnothing(topt) ? sopt : topt
+
+    single = select_option(sopts.single, topts.single)
+    allow_errors = select_option(sopts.allow_errors, topts.allow_errors)
+    proclist = select_option(sopts.proclist, topts.proclist)
     ThunkOptions(single,
                  proclist,
                  topts.time_util,

src/stream.jl

@@ -208,12 +208,10 @@ function initialize_streaming!(self_streams, spec, task)
         end
         output_buffer = get(spec.options, :stream_output_buffer, ProcessRingBuffer)
         stream = Stream{T,output_buffer}(output_buffer_amount)
-        spec.options = NamedTuple(filter(opt -> opt[1] != :stream_output_buffer &&
-                                                opt[1] != :stream_output_buffer_amount,
-                                         Base.pairs(spec.options)))
         self_streams[task.uid] = stream
 
-        spec.f = StreamingFunction(spec.f, stream)
+        max_evals = get(spec.options, :stream_max_evals, -1)
+        spec.f = StreamingFunction(spec.f, stream, max_evals)
         spec.options = merge(spec.options, (;occupancy=Dict(Any=>0)))
 
         # Register Stream globally
@@ -256,6 +254,7 @@ const STREAM_THUNK_ID = TaskLocalValue{Int}(()->0)
 struct StreamingFunction{F, S}
     f::F
     stream::S
+    max_evals::Int
 end
 chunktype(sf::StreamingFunction{F}) where F = F
 function (sf::StreamingFunction)(args...; kwargs...)
@@ -319,14 +318,17 @@ end
 function stream!(sf::StreamingFunction, uid,
                  args::Tuple, kwarg_names::Tuple, kwarg_values::Tuple)
     f = move(thunk_processor(), sf.f)
-    while true
+    counter = 0
+
+    while sf.max_evals < 0 || counter < sf.max_evals
         # Get values from Stream args/kwargs
         stream_args = _stream_take_values!(args, uid)
         stream_kwarg_values = _stream_take_values!(kwarg_values, uid)
         stream_kwargs = _stream_namedtuple(kwarg_names, stream_kwarg_values)
 
         # Run a single cycle of f
         stream_result = f(stream_args...; stream_kwargs...)
+        counter += 1
 
         # Exit streaming on graceful request
         if stream_result isa FinishStream
@@ -412,7 +414,8 @@ function finalize_streaming!(tasks::Vector{Pair{DTaskSpec,DTask}}, self_streams)
 
         # Filter out all streaming options
         to_filter = (:stream_input_buffer, :stream_input_buffer_amount,
-                     :stream_output_buffer, :stream_output_buffer_amount)
+                     :stream_output_buffer, :stream_output_buffer_amount,
+                     :stream_max_evals)
         spec.options = NamedTuple(filter(opt -> !(opt[1] in to_filter),
                                          Base.pairs(spec.options)))
         if haskey(spec.options, :propagates)

test/streaming.jl

@@ -5,6 +5,7 @@
     end
     return x
 end
+
 function catch_interrupt(f)
     try
         f()
@@ -17,6 +18,7 @@ function catch_interrupt(f)
         rethrow(err)
     end
 end
+
 function test_finishes(f, message::String; ignore_timeout=false)
     t = @eval Threads.@spawn @testset $message catch_interrupt($f)
     if timedwait(()->istaskdone(t), 10) == :timed_out
@@ -29,6 +31,7 @@ function test_finishes(f, message::String; ignore_timeout=false)
     end
     return true
 end
+
 @testset "Basics" begin
     @test test_finishes("Single task") do
         local x
@@ -50,6 +53,21 @@ end
         fetch(x)
     end
 
+    @test test_finishes("Max evaluations") do
+        counter = 0
+        function incrementer()
+            counter += 1
+        end
+
+        x = Dagger.with_options(; stream_max_evals=10) do
+            Dagger.spawn_streaming() do
+                Dagger.@spawn incrementer()
+            end
+        end
+        wait(x)
+        @test counter == 10
+    end
+
     @test test_finishes("Two tasks (sequential)") do
         local x, y
         @warn "\n\n\nStart streaming\n\n\n"