fix(map): stop queued parallel iterations after failure

Schedule bounded parallel map work lazily so queued iterations do not start after an aborting failure.

Also preserve the original execution error when final output computation raises during unwind, and keep aborting FailCog exceptions bubbling through system cogs that wrap nested scopes.

Agent: pi
Co-authored-by: ChatGPT <codex@openai.com>

Author: DerekStride

lib/roast/cog.rb

@@ -87,9 +87,9 @@ def run!(barrier, config, input_context, executor_scope_value, executor_scope_in
       rescue ControlFlow::FailCog => e
         # TODO: do something with the message passed to fail!
         @failed = true
-        # TODO: better / cleaner handling in the workflow execution manager for a workflow failure
-        #   just re-raising this exception for now
-        raise e if config.abort_on_failure?
+        # System cogs are wrappers around nested execution scopes. If a nested scope raises an aborting
+        # FailCog, keep bubbling it instead of converting it into a local, non-aborting system cog failure.
+        raise e if config.abort_on_failure? || is_a?(SystemCog)
       rescue ControlFlow::Next, ControlFlow::Break => e
         @skipped = true
         raise e

lib/roast/execution_manager.rb

@@ -92,26 +92,36 @@ def run!
       Sync do |sync_task|
         sync_task.annotate("ExecutionManager #{@scope}")
         TaskContext.begin_execution_manager(self)
-        @cog_stack.each do |cog|
-          cog_config = @config_manager.config_for(cog.class, cog.name)
-          cog_task = cog.run!(
-            @barrier,
-            cog_config.deep_dup,
-            cog_input_context,
-            @scope_value.deep_dup,
-            @scope_index,
-          )
-          cog_task.wait unless cog_config.async?
+        execution_error = nil
+        begin
+          @cog_stack.each do |cog|
+            cog_config = @config_manager.config_for(cog.class, cog.name)
+            cog_task = cog.run!(
+              @barrier,
+              cog_config.deep_dup,
+              cog_input_context,
+              @scope_value.deep_dup,
+              @scope_index,
+            )
+            cog_task.wait unless cog_config.async?
+          end
+          # Wait on the tasks in their completion order, so that an exception in a task will be raised as soon as it occurs
+          # noinspection RubyArgCount
+          @barrier.wait { |task| wait_for_task_with_exception_handling(task) }
+          compute_final_output # eagerly compute the final output (so it, too, can 'break!' subsequent executions in a loop)
+        rescue StandardError => e
+          execution_error = e
+          raise
+        ensure
+          @barrier.stop
+          begin
+            compute_final_output
+          rescue StandardError
+            raise unless execution_error
+          end
+          TaskContext.end
+          @running = false
         end
-        # Wait on the tasks in their completion order, so that an exception in a task will be raised as soon as it occurs
-        # noinspection RubyArgCount
-        @barrier.wait { |task| wait_for_task_with_exception_handling(task) }
-        compute_final_output # eagerly compute the final output (so it, too, can 'break!' subsequent executions in a loop)
-      ensure
-        @barrier.stop
-        compute_final_output
-        TaskContext.end
-        @running = false
       end
     end
 

lib/roast/system_cogs/map.rb

@@ -305,10 +305,17 @@ def execute_map_in_series(run, input)
         #: (Symbol, Map::Input, Integer?) -> Output
         def execute_map_in_parallel(run, input, max_parallel_tasks)
           barrier = Async::Barrier.new
-          semaphore = Async::Semaphore.new(max_parallel_tasks, parent: barrier) if max_parallel_tasks.present?
-          ems = {}
-          input.items.map.with_index do |item, index|
-            (semaphore || barrier).async(finished: false) do |task|
+          ems = Array.new(input.items.length)
+          next_index_to_schedule = 0
+
+          schedule_iteration = lambda do
+            next if next_index_to_schedule >= input.items.length
+
+            item = input.items[next_index_to_schedule]
+            index = next_index_to_schedule
+            next_index_to_schedule += 1
+
+            barrier.async(finished: false) do |task|
               task.annotate("Map Invocation #{index + input.initial_index}")
               ems[index] = em = create_execution_manager_for_map_item(run, item, index + input.initial_index)
               em.prepare!
@@ -317,12 +324,16 @@ def execute_map_in_parallel(run, input, max_parallel_tasks)
               # TODO: do something with the message passed to next!
               # proceed to next iteration
             end
-          end #: Array[Async::Task]
+          end
+
+          initial_parallelism = max_parallel_tasks || input.items.length
+          [initial_parallelism, input.items.length].min.times { schedule_iteration.call }
 
           # Wait on the tasks in their completion order, so that an exception in a task will be raised as soon as it occurs
           # noinspection RubyArgCount
           barrier.wait do |task|
             task.wait
+            schedule_iteration.call
           rescue ControlFlow::Break
             # TODO: do something with the message passed to break!
             barrier.stop
@@ -331,7 +342,7 @@ def execute_map_in_parallel(run, input, max_parallel_tasks)
             raise e
           end
 
-          Output.new((0...input.items.length).map { |idx| ems[idx] })
+          Output.new(ems)
         ensure
           # noinspection RubyRedundantSafeNavigation
           barrier&.stop

test/roast/system_cogs/map_test.rb

@@ -395,6 +395,50 @@ def build_manager(execution_procs, scope: nil, scope_value: nil, scope_index: 0)
         assert_equal "B", results[1].value
         assert_equal "C", results[2].value
       end
+
+      test "parallel map aborts pending iterations when an iteration fails and abort_on_failure is enabled" do
+        started_items = []
+
+        config_proc = proc do
+          map(:my_map) { parallel 2 }
+          test_cog { abort_on_failure! }
+        end
+        config_manager = ConfigManager.new(@registry, [config_proc])
+        config_manager.prepare!
+
+        exec_procs = {
+          nil => [proc {
+            map(:my_map, run: :parallel_failure) { ["slow", "fail", "later_1", "later_2"] }
+            outputs { collect(map(:my_map)) }
+          }],
+          parallel_failure: [proc {
+            test_cog(:step) do |_input, scope, _index|
+              started_items << scope
+
+              case scope
+              when "slow"
+                sleep 0.2
+                scope.upcase
+              when "fail"
+                raise ControlFlow::FailCog, "boom"
+              else
+                scope.upcase
+              end
+            end
+          }],
+        }
+        manager = ExecutionManager.new(
+          @registry, config_manager, exec_procs, @workflow_context
+        )
+        manager.prepare!
+
+        error = assert_raises(ControlFlow::FailCog) do
+          manager.run!
+        end
+
+        assert_equal "boom", error.message
+        assert_equal ["fail", "slow"], started_items.sort
+      end
     end
   end
 end