nv-ingest 2025.8.14.dev20250814__py3-none-any.whl → 2025.8.15.dev20250815__py3-none-any.whl

nv_ingest/framework/orchestration/ray/stages/mutate/image_filter.py

@@ -3,12 +3,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 
 import ray
 
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
 from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 from nv_ingest_api.internal.mutate.filter import filter_images_internal
 from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_task_by_type
 from nv_ingest_api.internal.primitives.tracing.tagging import traceable
@@ -31,18 +32,19 @@ class ImageFilterStage(RayActorStage):
       3. Updates the message payload with the filtered DataFrame.
     """
 
-    def __init__(self, config: ImageFilterSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: ImageFilterSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         try:
             self.validated_config = config
-            logger.info("ImageFilterStage configuration validated successfully.")
+            logger.debug("ImageFilterStage configuration validated successfully.")
         except Exception as e:
             logger.exception(f"Error validating Image Filter config: {e}")
             raise
 
-    @traceable("image_filter")
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     @filter_by_task(required_tasks=["filter"])
-    @nv_ingest_node_failure_try_except(annotation_id="image_filter", raise_on_failure=False)
     def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
         """
         Process the control message by filtering images.
@@ -57,7 +59,7 @@ class ImageFilterStage(RayActorStage):
         IngestControlMessage
             The updated message with filtered images in the payload.
         """
-        logger.info("ImageFilterStage.on_data: Starting image filtering process.")
+        logger.debug("ImageFilterStage.on_data: Starting image filtering process.")
 
         # Extract the DataFrame payload.
         df_ledger = control_message.payload()
@@ -76,7 +78,7 @@ class ImageFilterStage(RayActorStage):
             mutate_config=self.validated_config,
             execution_trace_log=None,
         )
-        logger.info("Image filtering completed. Resulting DataFrame has %d rows.", len(new_df))
+        logger.debug("Image filtering completed. Resulting DataFrame has %d rows.", len(new_df))
 
         # Update the message payload with the filtered DataFrame.
         control_message.payload(new_df)

nv_ingest/framework/orchestration/ray/stages/sinks/default_drain.py

@@ -2,7 +2,7 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 import ray
 
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_sink_stage_base import RayActorSinkStage
@@ -12,13 +12,13 @@ from nv_ingest_api.util.exception_handlers.decorators import nv_ingest_node_fail
 
 @ray.remote
 class DefaultDrainSink(RayActorSinkStage):
-    def __init__(self, config: Any) -> None:
-        super().__init__(config, log_to_stdout=False)
+    def __init__(self, config: Any, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, log_to_stdout=False, stage_name=stage_name)
 
         self._last_sunk_count = 0
         self._sunk_count = 0
 
-    @nv_ingest_node_failure_try_except(annotation_id="drain_sink", raise_on_failure=False)
+    @nv_ingest_node_failure_try_except()
     def on_data(self, message: IngestControlMessage) -> IngestControlMessage:
         self._sunk_count += 1
 

nv_ingest/framework/orchestration/ray/stages/sinks/message_broker_task_sink.py

@@ -14,6 +14,8 @@ from nv_ingest_api.internal.primitives.tracing.logging import annotate_cm
 from nv_ingest_api.util.message_brokers.simple_message_broker import SimpleClient
 from nv_ingest_api.util.service_clients.redis.redis_client import RedisClient
 
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
+
 logger = logging.getLogger(__name__)
 
 
@@ -75,8 +77,8 @@ class MessageBrokerTaskSinkConfig(BaseModel):
 
 @ray.remote
 class MessageBrokerTaskSinkStage(RayActorStage):
-    def __init__(self, config: MessageBrokerTaskSinkConfig) -> None:
-        super().__init__(config, log_to_stdout=False)
+    def __init__(self, config: MessageBrokerTaskSinkConfig, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, log_to_stdout=False, stage_name=stage_name)
 
         self.config: MessageBrokerTaskSinkConfig
 
@@ -224,6 +226,7 @@ class MessageBrokerTaskSinkStage(RayActorStage):
 
     # --- Public API Methods for message broker sink ---
 
+    @udf_intercept_hook()
     def on_data(self, control_message: Any) -> Any:
         """
         Processes the control message and pushes the resulting JSON payloads to the broker.

nv_ingest/framework/orchestration/ray/stages/sources/message_broker_task_source.py

@@ -3,9 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-import multiprocessing
 import uuid
-import socket
 from typing import Optional, Literal, Dict, Any, Union
 
 import ray
@@ -13,6 +11,7 @@ import json
 import copy
 import threading
 import time
+import random
 from datetime import datetime
 
 import pandas as pd
@@ -102,11 +101,11 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
     """
 
     # Use the updated config type hint
-    def __init__(self, config: MessageBrokerTaskSourceConfig) -> None:
-        super().__init__(config, log_to_stdout=False)
-        self.config: MessageBrokerTaskSourceConfig  # Add type hint for self.config
+    def __init__(self, config: MessageBrokerTaskSourceConfig, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, log_to_stdout=False, stage_name=stage_name)
+        self.config: MessageBrokerTaskSourceConfig  # Add a type hint for self.config
         self._logger.debug(
-            "Initializing MessageBrokerTaskSourceStage with config: %s", config.dict()
+            "Initializing MessageBrokerTaskSourceStage with config: %s", config.model_dump()
         )  # Log validated config
 
         # Access validated configuration directly via self.config
@@ -126,13 +125,18 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         self._pause_event = threading.Event()
         self._pause_event.set()  # Initially not paused
 
+        # Backoff state for graceful retries when broker is unavailable
+        self._fetch_failure_count: int = 0
+        self._current_backoff_sleep: float = 0.0
+        self._last_backoff_log_time: float = 0.0
+
         self._logger.debug("MessageBrokerTaskSourceStage initialized. Task queue: %s", self.task_queue)
 
     # --- Private helper methods ---
     def _create_client(self):
         # Access broker config via self.config.broker_client
         broker_config = self.config.broker_client
-        self._logger.info("Creating client of type: %s", broker_config.client_type)
+        self._logger.debug("Creating client of type: %s", broker_config.client_type)
 
         if broker_config.client_type == "redis":
             client = RedisClient(
@@ -265,6 +269,9 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
             job = self.client.fetch_message(self.task_queue, timeout)
             if job is None:
                 self._logger.debug("No message received from '%s'", self.task_queue)
+                # Do not treat normal empty polls as failures
+                self._fetch_failure_count = 0
+                self._current_backoff_sleep = 0.0
                 return None
             self._logger.debug("Received message type: %s", type(job))
             if isinstance(job, BaseModel):
@@ -277,12 +284,46 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
                     return None
                 job = json.loads(job.response)
             self._logger.debug("Successfully fetched message with job_id: %s", job.get("job_id", "unknown"))
+            # Success: reset backoff state
+            self._fetch_failure_count = 0
+            self._current_backoff_sleep = 0.0
             return job
         except TimeoutError:
             self._logger.debug("Timeout waiting for message")
+            # Timeout is not a connectivity failure; do not escalate backoff
             return None
         except Exception as err:
-            self._logger.exception("Error during message fetching: %s", err)
+            # Connectivity or other fetch issue: apply graceful backoff and avoid stacktrace spam
+            self._fetch_failure_count += 1
+
+            # Compute exponential backoff with jitter, capped by configured max_backoff
+            try:
+                max_backoff = getattr(self.config.broker_client, "max_backoff", 5.0)
+            except Exception:
+                max_backoff = 5.0
+            # Start from 0.5s, double each failure
+            base = 0.5
+            backoff_no_jitter = min(max_backoff, base * (2 ** (self._fetch_failure_count - 1)))
+            jitter = random.uniform(0, backoff_no_jitter * 0.2)
+            self._current_backoff_sleep = backoff_no_jitter + jitter
+
+            now = time.time()
+            # Throttle warning logs to at most once per 5 seconds to avoid spam
+            if now - self._last_backoff_log_time >= 5.0:
+                self._logger.warning(
+                    "Broker fetch failed (%d consecutive failures). Backing off for %.2fs. Error: %s",
+                    self._fetch_failure_count,
+                    self._current_backoff_sleep,
+                    err,
+                )
+                self._last_backoff_log_time = now
+            else:
+                self._logger.debug(
+                    "Broker fetch failed (%d). Backoff %.2fs. Error: %s",
+                    self._fetch_failure_count,
+                    self._current_backoff_sleep,
+                    err,
+                )
             return None
 
     def _read_input(self) -> any:
@@ -293,8 +334,17 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         self._logger.debug("read_input: calling _fetch_message()")
         job = self._fetch_message(timeout=100)
         if job is None:
-            self._logger.debug("read_input: No job received, sleeping for poll_interval: %s", self.config.poll_interval)
-            time.sleep(self.config.poll_interval)
+            # Sleep for either the configured poll interval or the current backoff, whichever is larger
+            sleep_time = max(self.config.poll_interval, getattr(self, "_current_backoff_sleep", 0.0))
+            self._logger.debug(
+                "read_input: No job received; sleeping %.2fs (poll_interval=%.2fs, backoff=%.2fs)",
+                sleep_time,
+                self.config.poll_interval,
+                getattr(self, "_current_backoff_sleep", 0.0),
+            )
+            time.sleep(sleep_time)
+            # Reset one-shot backoff so that repeated failures recompute progressively
+            self._current_backoff_sleep = 0.0
 
             return None
 
@@ -314,7 +364,7 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         This loop fetches messages from the broker and writes them to the output queue,
         but blocks on the pause event when the stage is paused.
         """
-        self._logger.info("Processing loop started")
+        self._logger.debug("Processing loop started")
         iteration = 0
         while self._running:
             iteration += 1
@@ -381,25 +431,25 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
                 self._active_processing = False
                 self._shutdown_signal_complete = True
 
-        self._logger.info("Processing loop ending")
+        self._logger.debug("Processing loop ending")
 
     @ray.method(num_returns=1)
     def start(self) -> bool:
         if self._running:
-            self._logger.info("Start called but stage is already running.")
+            self._logger.warning("Start called but stage is already running.")
             return False
         self._running = True
         self.start_time = time.time()
         self._message_count = 0
-        self._logger.info("Starting processing loop thread.")
+        self._logger.debug("Starting processing loop thread.")
         threading.Thread(target=self._processing_loop, daemon=True).start()
-        self._logger.info("MessageBrokerTaskSourceStage started.")
+        self._logger.debug("MessageBrokerTaskSourceStage started.")
         return True
 
     @ray.method(num_returns=1)
     def stop(self) -> bool:
         self._running = False
-        self._logger.info("Stop called on MessageBrokerTaskSourceStage")
+        self._logger.debug("Stop called on MessageBrokerTaskSourceStage")
         return True
 
     @ray.method(num_returns=1)
@@ -425,7 +475,7 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
     @ray.method(num_returns=1)
     def set_output_queue(self, queue_handle: any) -> bool:
         self.output_queue = queue_handle
-        self._logger.info("Output queue set: %s", queue_handle)
+        self._logger.debug("Output queue set: %s", queue_handle)
         return True
 
     @ray.method(num_returns=1)
@@ -440,7 +490,7 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
             True after the stage is paused.
         """
         self._pause_event.clear()
-        self._logger.info("Stage paused.")
+        self._logger.debug("Stage paused.")
 
         return True
 
@@ -456,7 +506,7 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
             True after the stage is resumed.
         """
         self._pause_event.set()
-        self._logger.info("Stage resumed.")
+        self._logger.debug("Stage resumed.")
         return True
 
     @ray.method(num_returns=1)
@@ -466,49 +516,9 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         This method pauses the stage, waits for any current processing to finish,
         replaces the output queue, and then resumes the stage.
         """
-        self._logger.info("Swapping output queue: pausing stage first.")
+        self._logger.debug("Swapping output queue: pausing stage first.")
         self.pause()
         self.set_output_queue(new_queue)
-        self._logger.info("Output queue swapped. Resuming stage.")
+        self._logger.debug("Output queue swapped. Resuming stage.")
         self.resume()
         return True
-
-
-def start_simple_message_broker(broker_client: dict) -> multiprocessing.Process:
-    """
-    Starts a SimpleMessageBroker server in a separate process.
-
-    Parameters
-    ----------
-    broker_client : dict
-        Broker configuration. Expected keys include:
-          - "port": the port to bind the server to,
-          - "broker_params": optionally including "max_queue_size",
-          - and any other parameters required by SimpleMessageBroker.
-
-    Returns
-    -------
-    multiprocessing.Process
-        The process running the SimpleMessageBroker server.
-    """
-
-    def broker_server():
-        from nv_ingest_api.util.message_brokers.simple_message_broker.broker import SimpleMessageBroker
-
-        # Use max_queue_size from broker_params or default to 10000.
-        broker_params = broker_client.get("broker_params", {})
-        max_queue_size = broker_params.get("max_queue_size", 10000)
-        server_host = broker_client.get("host", "0.0.0.0")
-        server_port = broker_client.get("port", 7671)
-        # Optionally, set socket options here for reuse.
-        server = SimpleMessageBroker(server_host, server_port, max_queue_size)
-        # Enable address reuse on the server socket.
-        server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        server.serve_forever()
-
-    p = multiprocessing.Process(target=broker_server)
-    p.daemon = False
-    p.start()
-    logger.info(f"Started SimpleMessageBroker server in separate process on port {broker_client['port']}")
-
-    return p

nv_ingest/framework/orchestration/ray/stages/storage/image_storage.py

@@ -3,13 +3,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 
 import pandas as pd
 import ray
 
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
 from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 from nv_ingest_api.internal.enums.common import ContentTypeEnum
 from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_task_by_type
 from nv_ingest_api.internal.primitives.tracing.tagging import traceable
@@ -31,8 +32,8 @@ class ImageStorageStage(RayActorStage):
     payload and updates the control message accordingly.
     """
 
-    def __init__(self, config: ImageStorageModuleSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: ImageStorageModuleSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         try:
             self.validated_config = config
             logger.info("ImageStorageStage configuration validated successfully.")
@@ -40,9 +41,10 @@ class ImageStorageStage(RayActorStage):
             logger.exception("Error validating image storage config")
             raise e
 
-    @traceable("image_storage")
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     @filter_by_task(required_tasks=["store"])
-    @nv_ingest_node_failure_try_except(annotation_id="image_storage", raise_on_failure=False)
     def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
         """
         Process the control message by storing images or structured content.

nv_ingest/framework/orchestration/ray/stages/storage/store_embeddings.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+from typing import Optional
 
 import ray
 
@@ -16,6 +17,8 @@ from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
 
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
+
 logger = logging.getLogger(__name__)
 
 
@@ -30,8 +33,8 @@ class EmbeddingStorageStage(RayActorStage):
       3. Updates the message payload with the stored embeddings DataFrame.
     """
 
-    def __init__(self, config: EmbeddingStorageSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: EmbeddingStorageSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         try:
             self.validated_config = config
             logger.info("EmbeddingStorageStage configuration validated successfully.")
@@ -39,9 +42,10 @@ class EmbeddingStorageStage(RayActorStage):
             logger.exception(f"Error validating Embedding Storage config: {e}")
             raise
 
-    @traceable("embedding_storage")
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     @filter_by_task(required_tasks=["store_embedding"])
-    @nv_ingest_node_failure_try_except(annotation_id="embedding_storage", raise_on_failure=False)
     def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
         """
         Process the control message by storing embeddings.

nv_ingest/framework/orchestration/ray/stages/telemetry/job_counter.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import Any
+from typing import Any, Optional
 from pydantic import BaseModel
 import ray
 
@@ -14,6 +14,8 @@ from nv_ingest.framework.util.telemetry.global_stats import GlobalStats
 from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
 
 # Import the JobCounter schema and global stats singleton.
 
@@ -30,15 +32,17 @@ class JobCounterStage(RayActorStage):
     statistic each time it processes a message.
     """
 
-    def __init__(self, config: BaseModel) -> None:
+    def __init__(self, config: BaseModel, stage_name: Optional[str] = None) -> None:
         # Ensure base attributes (e.g. self._running) are initialized.
-        super().__init__(config)
+        super().__init__(config, stage_name=stage_name)
         # The validated config should be a JobCounterSchema instance.
         self.validated_config: JobCounterSchema = config
         # Obtain the global stats' singleton.
         self.stats = GlobalStats.get_instance()
 
-    @nv_ingest_node_failure_try_except(annotation_id="job_counter", raise_on_failure=False)
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     async def on_data(self, message: Any) -> Any:
         """
         Process an incoming IngestControlMessage by counting jobs.

nv_ingest/framework/orchestration/ray/stages/telemetry/otel_tracer.py

@@ -24,6 +24,7 @@ from nv_ingest_api.util.exception_handlers.decorators import nv_ingest_node_fail
 
 from nv_ingest_api.internal.primitives.tracing.logging import TaskResultStatus
 from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 
 
 @ray.remote
@@ -35,8 +36,8 @@ class OpenTelemetryTracerStage(RayActorStage):
     It creates spans for tasks and exports them to a configured OpenTelemetry endpoint.
     """
 
-    def __init__(self, config: OpenTelemetryTracerSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: OpenTelemetryTracerSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
 
         # self._logger.info(f"[Telemetry] Initializing OpenTelemetry tracer stage with config: {config}")
 
@@ -81,7 +82,7 @@ class OpenTelemetryTracerStage(RayActorStage):
         parent_ctx = trace.set_span_in_context(NonRecordingSpan(span_context))
         parent_span = self.tracer.start_span(str(job_id), context=parent_ctx, start_time=start_time)
 
-        event_count = create_span_with_timestamps(self.tracer, parent_span, message)
+        event_count = create_span_with_timestamps(self.tracer, parent_span, message, self._logger)
 
         if message.has_metadata("cm_failed") and message.get_metadata("cm_failed"):
             parent_span.set_status(Status(StatusCode.ERROR))
@@ -96,7 +97,8 @@ class OpenTelemetryTracerStage(RayActorStage):
 
         self._logger.debug(f"[Telemetry] Exported spans for message {job_id} with {event_count} total events.")
 
-    @nv_ingest_node_failure_try_except(annotation_id="otel_tracer", raise_on_failure=False)
+    @nv_ingest_node_failure_try_except()
+    @udf_intercept_hook()
     def on_data(self, control_message: IngestControlMessage) -> Optional[Any]:
         try:
             do_trace_tagging = bool(control_message.get_metadata("config::add_trace_tagging"))
@@ -160,7 +162,7 @@ def extract_annotated_task_results(message):
     return task_results
 
 
-def create_span_with_timestamps(tracer, parent_span, message) -> int:
+def create_span_with_timestamps(tracer, parent_span, message, logger) -> int:
     timestamps = extract_timestamps_from_message(message)
     task_results = extract_annotated_task_results(message)
 
@@ -175,8 +177,16 @@ def create_span_with_timestamps(tracer, parent_span, message) -> int:
         if not subtask:
             span = tracer.start_span(main_task, context=child_ctx, start_time=ts_entry)
         else:
-            subtask_ctx = trace.set_span_in_context(ctx_store[main_task][0])
-            span = tracer.start_span(subtask, context=subtask_ctx, start_time=ts_entry)
+            # Check if parent context exists, otherwise create standalone span with warning
+            if main_task in ctx_store:
+                subtask_ctx = trace.set_span_in_context(ctx_store[main_task][0])
+                span = tracer.start_span(subtask, context=subtask_ctx, start_time=ts_entry)
+            else:
+                logger.warning(
+                    f"Missing parent context for subtask '{subtask}'"
+                    f" (expected parent: '{main_task}'). Creating standalone span."
+                )
+                span = tracer.start_span(f"{main_task}::{subtask}", context=child_ctx, start_time=ts_entry)
 
         span.add_event("entry", timestamp=ts_entry)
         span.add_event("exit", timestamp=ts_exit)

nv_ingest/framework/orchestration/ray/stages/transforms/image_caption.py

@@ -4,12 +4,13 @@
 
 import logging
 import pprint
-from typing import Any
+from typing import Any, Optional
 
 import ray
 
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
 from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type
 from nv_ingest_api.internal.primitives.tracing.tagging import traceable
 from nv_ingest_api.internal.schemas.transform.transform_image_caption_schema import ImageCaptionExtractionSchema
@@ -31,8 +32,8 @@ class ImageCaptionTransformStage(RayActorStage):
     are stored in the control message.
     """
 
-    def __init__(self, config: ImageCaptionExtractionSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: ImageCaptionExtractionSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         try:
             self.validated_config = config
             logger.info("ImageCaptionTransformStage configuration validated.")
@@ -40,9 +41,10 @@ class ImageCaptionTransformStage(RayActorStage):
             logger.exception("Error validating caption extraction config")
             raise e
 
-    @traceable("image_captioning")
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     @filter_by_task(required_tasks=["caption"])
-    @nv_ingest_node_failure_try_except(annotation_id="image_captioning", raise_on_failure=False)
     def on_data(self, control_message: Any) -> Any:
         """
         Process the control message by extracting image captions.

nv_ingest/framework/orchestration/ray/stages/transforms/text_embed.py

@@ -2,12 +2,10 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
 import pprint
-from typing import Any
+from typing import Optional
 import ray
 
-# Assume these imports come from your project:
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
 from nv_ingest.framework.util.flow_control import filter_by_task
 from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type, IngestControlMessage
@@ -18,7 +16,7 @@ from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
 
-logger = logging.getLogger(__name__)
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 
 
 @ray.remote
@@ -31,19 +29,20 @@ class TextEmbeddingTransformStage(RayActorStage):
     trace or extraction metadata is added.
     """
 
-    def __init__(self, config: TextEmbeddingSchema) -> None:
-        super().__init__(config, log_to_stdout=False)
+    def __init__(self, config: TextEmbeddingSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         try:
             self.validated_config = config
-            logger.info("TextEmbeddingTransformStage configuration validated successfully.")
+            self._logger.info("TextEmbeddingTransformStage configuration validated successfully.")
         except Exception as e:
-            logger.exception("Error validating text embedding extractor config")
-            raise e
+            self._logger.exception(f"Error validating text embedding config: {e}")
+            raise
 
-    @traceable("text_embedding")
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     @filter_by_task(required_tasks=["embed"])
-    @nv_ingest_node_failure_try_except(annotation_id="text_embedding", raise_on_failure=False)
-    def on_data(self, control_message: IngestControlMessage) -> Any:
+    def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
         """
         Process the control message by generating text embeddings.
 
@@ -59,11 +58,11 @@ class TextEmbeddingTransformStage(RayActorStage):
         """
         # Get the DataFrame payload.
         df_payload = control_message.payload()
-        logger.debug("TextEmbeddingTransformStage: Extracted payload with %d rows.", len(df_payload))
+        self._logger.debug("TextEmbeddingTransformStage: Extracted payload with %d rows.", len(df_payload))
 
         # Remove the "embed" task to obtain task-specific configuration.
         task_config = remove_task_by_type(control_message, "embed")
-        logger.debug("TextEmbeddingTransformStage: Task configuration extracted: %s", pprint.pformat(task_config))
+        self._logger.debug("TextEmbeddingTransformStage: Task configuration extracted: %s", pprint.pformat(task_config))
 
         # Call the text embedding extraction function.
         new_df, execution_trace_log = transform_create_text_embeddings_internal(