@@ -574,9 +574,6 @@ def clear_completed_children(self) -> None:
574574 """
575575 Clears all of this task's children which are complete.
576576 """
577- if self ._next_child_index != 0 :
578- raise ValueError ('Unsafe to call clear_completed_children unless _next_child_index == 0' )
579-
580577 child_indexes_to_remove = [i for (i , c ) in enumerate (self ._children ) if c .complete ] # capture
581578 if len (child_indexes_to_remove ) == 0 :
582579 return
@@ -588,6 +585,7 @@ def clear_completed_children(self) -> None:
588585 )
589586 self ._children = [c for c in self .children if not c .complete ]
590587 self ._num_children_complete = 0
588+ self ._next_child_index = 0
591589
592590 # Notify listeners last, which can run arbitrary code
593591 # NOTE: Call these listeners also inside the lock
@@ -674,18 +672,21 @@ def try_get_next_task_unit(self) -> Callable[[], None] | None:
674672 else :
675673 return None
676674 elif self .scheduling_style == SCHEDULING_STYLE_ROUND_ROBIN :
677- if self ._next_child_index == 0 :
678- # NOTE: Ignore known-slow operation that has no further obvious optimizations
679- with ignore_runtime_from_enclosing_warn_if_slow ():
680- schedule_check_result = self ._notify_did_schedule_all_children ()
681- if not isinstance (schedule_check_result , bool ):
682- return schedule_check_result
683675 any_incomplete_interactive_children = any (
684676 c .interactive and not c .complete
685677 for c in self .children
686678 ) # cache
687679 cur_child_index = self ._next_child_index
688680 while True :
681+ if cur_child_index == 0 :
682+ # NOTE: Ignore known-slow operation that has no further obvious optimizations
683+ with ignore_runtime_from_enclosing_warn_if_slow ():
684+ self ._notify_did_schedule_all_children ()
685+ # (self.children and self._next_child_index may have changed)
686+ if len (self .children ) == 0 :
687+ # Handle zero-children case in usual manner
688+ return self .try_get_next_task_unit ()
689+ assert cur_child_index == 0
689690 cur_child = self .children [cur_child_index ]
690691 if any_incomplete_interactive_children and not cur_child .interactive :
691692 # If there are any interactive priority children
@@ -708,37 +709,12 @@ def try_get_next_task_unit(self) -> Callable[[], None] | None:
708709 # as one of the crashed children.
709710 self .crash_reason = first_crashed_child .crash_reason
710711 return None
711- if cur_child_index == 0 :
712- # Temporarily force self._next_child_index to appear as 0
713- # because some listeners called by self._notify_did_schedule_all_children()
714- # may expect it to be exactly 0. Notably RootTask.did_schedule_all_children().
715- old_next_child_index = self ._next_child_index
716- self ._next_child_index = 0
717- try :
718- # NOTE: Ignore known-slow operation that has no further obvious optimizations
719- with ignore_runtime_from_enclosing_warn_if_slow ():
720- schedule_check_result = self ._notify_did_schedule_all_children ()
721- finally :
722- self ._next_child_index = old_next_child_index
723- if not isinstance (schedule_check_result , bool ):
724- return schedule_check_result
725- elif schedule_check_result == True :
726- # Invalidate self._next_child_index,
727- # because children may have changed
728- self ._next_child_index = 0
729712 else :
730713 raise ValueError ('Container task has an unknown scheduling style (%s).' % self .scheduling_style )
731714
732- def _notify_did_schedule_all_children (self ) -> bool | Callable [[], None ] | None :
715+ def _notify_did_schedule_all_children (self ) -> None :
733716 if hasattr (self , 'did_schedule_all_children' ):
734717 self .did_schedule_all_children () # type: ignore[attr-defined]
735- # (Children may have changed)
736- if len (self .children ) == 0 :
737- # Handle zero-children case in usual manner
738- return self .try_get_next_task_unit ()
739- return True # children may have changed
740- else :
741- return False # children did not change
742718
743719 @bg_affinity
744720 @capture_crashes_to_self
0 commit comments