nv-ingest 2025.7.8.dev20250708__py3-none-any.whl → 2025.7.9.dev20250709__py3-none-any.whl

nv_ingest/framework/orchestration/ray/primitives/pipeline_topology.py

@@ -6,7 +6,6 @@ import threading
 import logging
 import contextlib
 import time
-from collections import defaultdict
 from typing import List, Dict, Tuple, Any, Optional, Iterator, Set
 
 import ray
@@ -59,19 +58,15 @@ class PipelineTopology:
         self._edge_queues: Dict[str, Tuple[Any, int]] = {}  # Map: q_name -> (QueueHandle, Capacity)
         self._scaling_state: Dict[str, str] = {}  # Map: stage_name -> "Idle" | "Scaling Up" | "Scaling Down" | "Error"
         self._stage_memory_overhead: Dict[str, float] = {}  # Populated during build/config
-        self._pending_removal_actors: Dict[str, Set[Tuple[Any, str, float, ray.ObjectRef]]] = defaultdict(set)
+        self._actors_pending_removal: Set[Tuple[str, Any]] = set()
 
         # --- Operational State ---
         self._is_flushing: bool = False
 
         # --- Synchronization & Threading ---
         self._lock: threading.Lock = threading.Lock()
-        self._cleanup_thread: Optional[threading.Thread] = None
-        self._cleanup_thread_running: bool = False
-        self._stop_event: threading.Event = threading.Event()  # For interruptible sleep
-
-        logger.debug("PipelineTopology initialized.")
-        self._start_cleanup_thread()  # Start background cleanup on init
+        self._stop_cleanup = threading.Event()
+        self._cleanup_thread = None
 
     def __del__(self):
         """Ensure cleanup thread is stopped and internal actor references are released."""
@@ -79,7 +74,7 @@ class PipelineTopology:
 
         # Stop the background cleanup thread
         try:
-            self._stop_cleanup_thread()
+            self.stop_cleanup_thread()
         except Exception as e:
             logger.warning(f"Error stopping cleanup thread during __del__: {e}")
 
@@ -89,7 +84,7 @@ class PipelineTopology:
             self._edge_queues.clear()
             self._scaling_state.clear()
             self._stage_memory_overhead.clear()
-            self._pending_removal_actors.clear()
+            self._actors_pending_removal.clear()
             self._stages.clear()
             self._connections.clear()
         except Exception as e:
@@ -184,303 +179,65 @@ class PipelineTopology:
 
         return removed
 
-    def register_actors_pending_removal(self, registration_info: Dict[str, List[Tuple[Any, ray.ObjectRef]]]) -> None:
-        """
-        Registers actor handles that have been told to stop, along with their shutdown futures.
-        The topology's background thread will monitor these futures for completion.
-
-        Parameters
-        ----------
-        registration_info : Dict[str, List[Tuple[Any, ObjectRef]]]
-            Dictionary mapping stage names to a list of (actor_handle, shutdown_future) tuples.
-        """
-        added_count = 0
-        time_registered = time.time()
-        stages_updated = set()
-
+    def mark_actor_for_removal(self, stage_name: str, actor: Any) -> None:
+        """Marks an actor as pending removal, to be cleaned up by the background thread."""
         with self._lock:
-            all_known_stages = {s.name for s in self._stages}
-
-            for stage_name, actor_list in registration_info.items():
-                if stage_name not in all_known_stages:
+            self._actors_pending_removal.add((stage_name, actor))
+            logger.info(f"Marked actor {actor} from stage {stage_name} for removal.")
+
+    def start_cleanup_thread(self, interval: int = 5) -> None:
+        """Starts the background thread for periodic cleanup tasks."""
+        if self._cleanup_thread is None or not self._cleanup_thread.is_alive():
+            self._stop_cleanup.clear()
+            self._cleanup_thread = threading.Thread(target=self._cleanup_loop, args=(interval,), daemon=True)
+            self._cleanup_thread.start()
+            logger.info("Topology cleanup thread started.")
+
+    def stop_cleanup_thread(self) -> None:
+        """Stops the background cleanup thread."""
+        if self._cleanup_thread and self._cleanup_thread.is_alive():
+            self._stop_cleanup.set()
+            self._cleanup_thread.join(timeout=5)
+            logger.info("Topology cleanup thread stopped.")
+
+    def _cleanup_loop(self, interval: int) -> None:
+        """Periodically checks for and removes actors that have completed shutdown."""
+        while not self._stop_cleanup.is_set():
+            actors_to_remove_finally = []
+            if not self._actors_pending_removal:
+                time.sleep(interval)
+                continue
+
+            # Check the status of actors pending removal
+            # Create a copy for safe iteration, as the set might be modified elsewhere
+            pending_actors_copy = set()
+            with self._lock:
+                pending_actors_copy = set(self._actors_pending_removal)
+
+            for stage_name, actor in pending_actors_copy:
+                try:
+                    if ray.get(actor.is_shutdown_complete.remote()):
+                        actors_to_remove_finally.append((stage_name, actor))
+                except ray.exceptions.RayActorError:
                     logger.warning(
-                        f"[TopologyRegister] Received pending removal registration for unknown stage "
-                        f"'{stage_name}'. Skipping."
+                        f"Actor {actor} from stage {stage_name} is no longer available (RayActorError). "
+                        f"Assuming it has shut down and marking for removal."
                     )
-                    continue
-
-                stage_pending_set = self._pending_removal_actors[stage_name]
-
-                for actor_handle, shutdown_future in actor_list:
-                    if not actor_handle or not shutdown_future:
-                        logger.warning(
-                            f"[TopologyRegister-{stage_name}] "
-                            f"Received invalid (actor, future) in registration list. Skipping."
-                        )
-                        continue
-
-                    actor_id_str = str(actor_handle)
-                    actor_tuple = (actor_handle, actor_id_str, time_registered, shutdown_future)
-
-                    if actor_tuple not in stage_pending_set:
-                        stage_pending_set.add(actor_tuple)
-                        added_count += 1
-                        logger.debug(
-                            f"[TopologyRegister-{stage_name}] "
-                            f"Registered actor '{actor_id_str}' pending shutdown monitoring."
-                        )
-                    else:
-                        logger.debug(
-                            f"[TopologyRegister-{stage_name}] "
-                            f"Actor '{actor_id_str}' already registered pending removal."
-                        )
-
-                if actor_list:
-                    self._scaling_state[stage_name] = "Scaling Down Pending"
-                    stages_updated.add(stage_name)
-
-        if added_count > 0:
-            logger.debug(
-                f"[TopologyRegister] Registered {added_count} "
-                f"actors across {len(stages_updated)} stages pending removal."
-            )
-        elif registration_info:
-            logger.debug("[TopologyRegister] No new actors registered pending removal (likely duplicates).")
-
-    def _start_cleanup_thread(self) -> None:
-        """Starts the background thread for cleaning up terminated actors."""
-        with self._lock:  # Protect thread state modification
-            if self._cleanup_thread is not None and self._cleanup_thread.is_alive():
-                logger.warning("[TopologyCleanup] Cleanup thread already started.")
-                return
-
-            logger.info("[TopologyCleanup] Starting background cleanup thread...")
-            self._cleanup_thread_running = True
-            self._stop_event.clear()  # Ensure event is not set initially
-            self._cleanup_thread = threading.Thread(
-                target=self._cleanup_loop,
-                daemon=True,  # Allows program exit even if this thread hangs (though join tries)
-                name="TopologyActorCleanup",
-            )
-            self._cleanup_thread.start()
-
-    def _stop_cleanup_thread(self) -> None:
-        """Signals the background cleanup thread to stop and waits for it."""
-        if not self._cleanup_thread_running or self._cleanup_thread is None:
-            logger.debug("[TopologyCleanup] Cleanup thread not running or already stopped.")
-            return
-
-        with self._lock:  # Protect thread state read/write
-            if not self._cleanup_thread_running or self._cleanup_thread is None:
-                return  # Double check inside lock
-            logger.info("[TopologyCleanup] Stopping background cleanup thread...")
-            self._cleanup_thread_running = False
-            self._stop_event.set()  # Signal the loop to wake up and exit
-
-        # Wait for the thread to finish outside the lock
-        join_timeout = CLEANUP_INTERVAL_SECONDS + 5.0  # Give it time to finish last cycle
-        self._cleanup_thread.join(timeout=join_timeout)
-
-        if self._cleanup_thread.is_alive():
-            logger.warning(f"[TopologyCleanup] Cleanup thread did not exit gracefully after {join_timeout}s.")
-        else:
-            logger.info("[TopologyCleanup] Cleanup thread stopped and joined.")
-        self._cleanup_thread = None  # Clear thread object
-
-    @staticmethod
-    def _delayed_actor_release(self, actor_handle_to_release: Any, actor_id_str: str, delay_seconds: int = 60):
-        """
-        Holds a reference to an actor handle for a specified delay, then releases it.
-        This function is intended to be run in a daemon thread.
-
-        Note: this is a bit of a hack
-        """
-        logger.debug(f"[DelayedRelease-{actor_id_str}] Thread started. Holding actor reference for {delay_seconds}s.")
-        # The actor_handle_to_release is kept in scope by being a parameter to this function,
-        # and this function's frame existing for delay_seconds.
-        time.sleep(delay_seconds)
-        logger.info(
-            f"[DelayedRelease-{actor_id_str}] Delay complete. Releasing reference. Actor should now be GC'd by Ray "
-            f"if this was the last ref."
-        )
-        # When this function exits, actor_handle_to_release goes out of scope, dropping the reference.
-
-    def _cleanup_loop(self) -> None:
-        """
-        Background thread for periodically checking shutdown status of actors pending removal.
+                    actors_to_remove_finally.append((stage_name, actor))
+                except Exception as e:
+                    logger.error(f"Error checking shutdown status for actor {actor}: {e}", exc_info=True)
 
-        Actors are removed from the topology once their shutdown futures complete or they time out.
-        """
-        logger.info("[TopologyCleanupLoop] Cleanup thread started.")
-
-        while self._cleanup_thread_running:
-            cycle_start_time = time.time()
-            actors_removed_this_cycle = 0
-            processed_actor_ids_this_cycle = set()
-            actors_to_remove_from_pending: Dict[str, List[Tuple[Any, str, float, ray.ObjectRef]]] = defaultdict(list)
-            stages_potentially_idle: Set[str] = set()
-
-            try:
+            # Remove the fully shut-down actors from the topology
+            if actors_to_remove_finally:
                 with self._lock:
-                    if not self._cleanup_thread_running:
-                        logger.debug(
-                            "[TopologyCleanupLoop] " "Stop signal received after lock acquisition. Exiting loop."
-                        )
-                        break
-
-                    for stage_name in list(self._pending_removal_actors.keys()):
-                        pending_set = self._pending_removal_actors[stage_name]
-                        if not pending_set:
-                            continue
-
-                        pending_set_copy = pending_set.copy()
-
-                        for actor_tuple in pending_set_copy:
-                            actor_handle, actor_id_str, time_registered, shutdown_future = actor_tuple
-
-                            if actor_id_str in processed_actor_ids_this_cycle:
-                                continue
-
-                            remove_from_topology = False
-                            mark_for_pending_removal = False
-                            actor_status = "PENDING"
-
-                            # 1. Check for overall shutdown timeout
-                            if time.time() - time_registered > PENDING_SHUTDOWN_TIMEOUT_SECONDS:
-                                logger.warning(
-                                    f"[TopologyCleanupLoop-{stage_name}] Actor '{actor_id_str}' "
-                                    f"timed out after {PENDING_SHUTDOWN_TIMEOUT_SECONDS}s. Forcing removal."
-                                )
-                                remove_from_topology = True
-                                mark_for_pending_removal = True
-                                actor_status = "TIMEOUT"
-
-                            # 2. Otherwise, check if shutdown future completed
-                            if not remove_from_topology:
-                                try:
-                                    ready, _ = ray.wait([shutdown_future], timeout=PENDING_CHECK_ACTOR_METHOD_TIMEOUT)
-                                    if ready:
-                                        logger.debug(
-                                            f"[TopologyCleanupLoop-{stage_name}] "
-                                            f"Actor '{actor_id_str}' shutdown future completed. Marking for removal."
-                                        )
-                                        remove_from_topology = True
-                                        mark_for_pending_removal = True
-                                        actor_status = "COMPLETED"
-                                    else:
-                                        logger.debug(
-                                            f"[TopologyCleanupLoop-{stage_name}] "
-                                            f"Actor '{actor_id_str}' shutdown future still pending."
-                                        )
-                                        actor_status = "PENDING"
-                                except Exception as e:
-                                    logger.error(
-                                        f"[TopologyCleanupLoop-{stage_name}] "
-                                        f"Error checking shutdown future for actor '{actor_id_str}': {e}",
-                                        exc_info=False,
-                                    )
-                                    actor_status = "ERROR"
-
-                            # 3. Perform removal actions
-                            if remove_from_topology:
-                                logger.debug(
-                                    f"[TopologyCleanupLoop-{stage_name}] Removing actor '{actor_id_str}' "
-                                    f"from topology (Reason: {actor_status})."
-                                )
-                                removed_list = self.remove_actors_from_stage(stage_name, [actor_handle])
-                                if removed_list:
-                                    actors_removed_this_cycle += 1
-                                else:
-                                    logger.debug(
-                                        f"[TopologyCleanupLoop-{stage_name}] Actor '{actor_id_str}' "
-                                        f"was already removed from main list."
-                                    )
-
-                            if mark_for_pending_removal:
-                                actors_to_remove_from_pending[stage_name].append(actor_tuple)
-                                processed_actor_ids_this_cycle.add(actor_id_str)
-                                stages_potentially_idle.add(stage_name)
-
-                    # --- Update pending lists ---
-                    for stage_to_update, removal_list in actors_to_remove_from_pending.items():
-                        if stage_to_update in self._pending_removal_actors:
-                            current_pending_set = self._pending_removal_actors[stage_to_update]
-                            for removal_tuple in removal_list:  # removal_list contains actor_tuples
-                                # Extract actor_handle and actor_id_str from the tuple being removed
-                                actor_handle_to_delay, actor_id_str_to_delay, _, _ = removal_tuple
-
-                                if current_pending_set.discard(
-                                    removal_tuple
-                                ):  # If discard was successful (element was present)
-                                    logger.debug(
-                                        f"[TopologyCleanupLoop-{stage_to_update}] Actor tuple for "
-                                        f"'{actor_id_str_to_delay}' discarded from pending set."
-                                    )
-                                    try:
-                                        # This is a bit of a hack. For some reason Ray likes to cause exceptions on
-                                        # the actor when we let it auto GCS just after pushing to the output queue, and
-                                        # mysteriously lose control messages.
-                                        # This lets the shutdown future complete, but leaves the actor to be killed off
-                                        # by ray.actor_exit()
-                                        delay_thread = threading.Thread(
-                                            target=self._delayed_actor_release,
-                                            args=(actor_handle_to_delay, actor_id_str_to_delay, 60),  # 60s delay
-                                            daemon=True,
-                                        )
-                                        delay_thread.start()
-                                        logger.debug(
-                                            f"[TopologyCleanupLoop-{stage_to_update}] Started delayed release thread "
-                                            f"for '{actor_id_str_to_delay}'."
-                                        )
-                                    except Exception as e_thread:
-                                        logger.error(
-                                            f"[TopologyCleanupLoop-{stage_to_update}] Failed to start delayed release "
-                                            f"thread for '{actor_id_str_to_delay}': {e_thread}"
-                                        )
-
-                            # After processing all removals for this stage's list, check if the set is empty
-                            if not self._pending_removal_actors[stage_to_update]:
-                                logger.debug(
-                                    f"[TopologyCleanupLoop-{stage_to_update}] Pending set empty. Deleting key."
-                                )
-                                del self._pending_removal_actors[stage_to_update]
-
-                    # --- Update stage scaling states if pending list is empty ---
-                    stages_with_empty_pending = []
-                    stages_with_empty_pending = []
-                    for stage_to_check in stages_potentially_idle:
-                        if stage_to_check not in self._pending_removal_actors:
-                            stages_with_empty_pending.append(stage_to_check)
-                            if self._scaling_state.get(stage_to_check) == "Scaling Down Pending":
-                                logger.debug(  # Your original log level
-                                    f"[TopologyCleanupLoop-{stage_to_check}] All pending actors cleared. "
-                                    f"Setting scaling state to Idle."
-                                )
-                                self._scaling_state[stage_to_check] = "Idle"
-
-                    # --- Log cycle summary ---
-                    cycle_duration = time.time() - cycle_start_time
-                    if actors_removed_this_cycle > 0:
-                        logger.debug(
-                            f"[TopologyCleanupLoop] Cleanup cycle finished in {cycle_duration:.3f}s. "
-                            f"Removed {actors_removed_this_cycle} actors."
-                        )
-                    else:
-                        logger.debug(
-                            f"[TopologyCleanupLoop] Cleanup cycle finished in {cycle_duration:.3f}s. "
-                            f"No actors removed."
-                        )
-
-            except Exception as e:
-                logger.error(f"[TopologyCleanupLoop] Unhandled error in cleanup loop iteration: " f"{e}", exc_info=True)
-
-            # --- Wait until next cycle ---
-            woken_by_stop = self._stop_event.wait(timeout=CLEANUP_INTERVAL_SECONDS)
-            if woken_by_stop:
-                logger.info("[TopologyCleanupLoop] Stop event received during sleep. Exiting loop.")
-                break
-
-        logger.info("[TopologyCleanupLoop] Cleanup thread finished.")
+                    for stage_name, actor in actors_to_remove_finally:
+                        if (stage_name, actor) in self._actors_pending_removal:
+                            self._actors_pending_removal.remove((stage_name, actor))
+                        if actor in self._stage_actors.get(stage_name, []):
+                            self._stage_actors[stage_name].remove(actor)
+                            logger.info(f"Successfully removed actor {actor} from stage {stage_name} in topology.")
+
+            time.sleep(interval)
 
     def set_edge_queues(self, queues: Dict[str, Tuple[Any, int]]) -> None:
         """Sets the dictionary of edge queues."""
@@ -526,12 +283,17 @@ class PipelineTopology:
 
     # --- Accessor Methods (Read Operations - Use Lock, Return Copies) ---
 
+    def get_all_actors(self) -> List[Any]:
+        """Returns a list of all actors across all stages."""
+        with self._lock:
+            return [actor for actors in self._stage_actors.values() for actor in actors]
+
     def get_stages_info(self) -> List[StageInfo]:
         """Returns a copy of stage info with pending_shutdown flags updated."""
         with self._lock:
             updated_stages = []
             for stage in self._stages:
-                pending_shutdown = bool(self._pending_removal_actors.get(stage.name))
+                pending_shutdown = bool(self._actors_pending_removal)
                 # Make a shallow copy with updated pending_shutdown
                 stage_copy = StageInfo(
                     name=stage.name,