@@ -551,27 +551,41 @@ def _task_success(
551551 group_events : list [GroupValuesEvent | GroupSuccessEvent ] = []
552552 new_groups : list [Group ] = []
553553 new_streams : list [Stream ] = []
554- root_groups = self ._root_groups
555554 for group in task .groups :
556555 group_node = self ._group_nodes .get (group )
557556 if group_node :
558557 group_node .pending -= 1
559- if group in root_groups and not group_node .pending :
560- (
561- group_values_event ,
562- group_success_event ,
563- child_new_groups ,
564- child_new_streams ,
565- ) = self ._finish_group_success (group , group_node )
566- if group_values_event :
567- group_events .append (group_values_event )
568- group_events .append (group_success_event )
569- new_groups .extend (child_new_groups )
570- new_streams .extend (child_new_streams )
558+ self ._maybe_finish_group (
559+ group , group_node , group_events , new_groups , new_streams
560+ )
571561
572562 self ._start_new_work (new_groups , new_streams )
573563 return group_events
574564
565+ def _maybe_finish_group (
566+ self ,
567+ group : Group ,
568+ group_node : _GroupNode ,
569+ group_events : list [GroupValuesEvent | GroupSuccessEvent ],
570+ new_groups : list [Group ],
571+ new_streams : list [Stream ],
572+ ) -> None :
573+ """Finish a root group if it has no pending tasks left."""
574+ if group not in self ._root_groups :
575+ return
576+ if not group_node .pending :
577+ (
578+ group_values_event ,
579+ group_success_event ,
580+ child_new_groups ,
581+ child_new_streams ,
582+ ) = self ._finish_group_success (group , group_node )
583+ if group_values_event :
584+ group_events .append (group_values_event )
585+ group_events .append (group_success_event )
586+ new_groups .extend (child_new_groups )
587+ new_streams .extend (child_new_streams )
588+
575589 def _task_failure (self , graph_event : _TaskFailure ) -> Sequence [GroupFailureEvent ]:
576590 """Handle a task failure, removing all groups it belongs to."""
577591 task , error = graph_event
0 commit comments