nv-ingest 2025.8.4.dev20250804__py3-none-any.whl → 2025.12.10.dev20251210__py3-none-any.whl

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
@@ -17,6 +18,7 @@ from nv_ingest_api.internal.transform.caption_image import transform_image_creat
 from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
+from nv_ingest_api.util.logging.sanitize import sanitize_for_logging
 
 logger = logging.getLogger(__name__)
 
@@ -31,8 +33,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 +42,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.
@@ -65,7 +68,10 @@ class ImageCaptionTransformStage(RayActorStage):
 
         # Remove the "caption" task to obtain task-specific configuration.
         task_config = remove_task_by_type(control_message, "caption")
-        logger.debug("ImageCaptionTransformStage: Task configuration extracted: %s", pprint.pformat(task_config))
+        logger.debug(
+            "ImageCaptionTransformStage: Task configuration extracted: %s",
+            pprint.pformat(sanitize_for_logging(task_config)),
+        )
 
         # Call the caption extraction function.
         new_df = transform_image_create_vlm_caption_internal(

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
@@ -17,8 +15,9 @@ from nv_ingest_api.internal.transform.embed_text import transform_create_text_em
 from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
+from nv_ingest_api.util.logging.sanitize import sanitize_for_logging
 
-logger = logging.getLogger(__name__)
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 
 
 @ray.remote
@@ -31,19 +30,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.
 
@@ -57,25 +57,24 @@ class TextEmbeddingTransformStage(RayActorStage):
         IngestControlMessage
             The updated message with text embeddings and trace info added.
         """
-        logger.info("TextEmbeddingTransformStage.on_data: Starting text embedding transformation.")
-
         # 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(sanitize_for_logging(task_config)),
+        )
 
         # Call the text embedding extraction function.
         new_df, execution_trace_log = transform_create_text_embeddings_internal(
             df_payload, task_config=task_config, transform_config=self.validated_config
         )
-        logger.info("Text embedding transformation completed. New payload has %d rows.", len(new_df))
 
         # Update the control message payload.
         control_message.payload(new_df)
         # Annotate the message metadata with trace info.
         control_message.set_metadata("text_embedding_trace", execution_trace_log)
-        logger.info("Text embedding trace metadata added.")
         return control_message

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

@@ -3,18 +3,22 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-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_api.internal.primitives.ingest_control_message import remove_task_by_type, IngestControlMessage
+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
 from nv_ingest_api.internal.schemas.transform.transform_text_splitter_schema import TextSplitterSchema
 from nv_ingest_api.internal.transform.split_text import transform_text_split_and_tokenize_internal
 from nv_ingest_api.util.exception_handlers.decorators import (
     nv_ingest_node_failure_try_except,
 )
+from nv_ingest_api.util.logging.sanitize import sanitize_for_logging
+
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
 
 logger = logging.getLogger(__name__)
 
@@ -29,15 +33,16 @@ class TextSplitterStage(RayActorStage):
     and tokenization logic. The updated DataFrame is then set back into the message.
     """
 
-    def __init__(self, config: TextSplitterSchema) -> None:
-        super().__init__(config)
+    def __init__(self, config: TextSplitterSchema, stage_name: Optional[str] = None) -> None:
+        super().__init__(config, stage_name=stage_name)
         # Store the validated configuration (assumed to be an instance of TextSplitterSchema)
         self.validated_config: TextSplitterSchema = config
-        logger.info("TextSplitterStage initialized with config: %s", config)
+        logger.info("TextSplitterStage initialized with config: %s", sanitize_for_logging(config))
 
-    @traceable("text_splitter")
-    @filter_by_task(["split"])
-    @nv_ingest_node_failure_try_except(annotation_id="text_splitter", raise_on_failure=False)
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
+    @filter_by_task(required_tasks=["split"])
     def on_data(self, message: Any) -> Any:
         """
         Process an incoming IngestControlMessage by splitting and tokenizing its text.
@@ -59,7 +64,7 @@ class TextSplitterStage(RayActorStage):
 
         # Remove the "split" task to obtain task-specific configuration.
         task_config = remove_task_by_type(message, "split")
-        logger.debug("Extracted task config: %s", task_config)
+        logger.debug("Extracted task config: %s", sanitize_for_logging(task_config))
 
         # Transform the DataFrame (split text and tokenize).
         df_updated = transform_text_split_and_tokenize_internal(
@@ -68,11 +73,13 @@ class TextSplitterStage(RayActorStage):
             transform_config=self.validated_config,
             execution_trace_log=None,
         )
-        logger.info("TextSplitterStage.on_data: Transformation complete. Updated payload has %d rows.", len(df_updated))
+        logger.debug(
+            "TextSplitterStage.on_data: Transformation complete. Updated payload has %d rows.", len(df_updated)
+        )
 
         # Update the message payload.
         message.payload(df_updated)
-        logger.info("TextSplitterStage.on_data: Finished processing, returning updated message.")
+        logger.debug("TextSplitterStage.on_data: Finished processing, returning updated message.")
 
         return message
 
@@ -101,7 +108,7 @@ def text_splitter_fn(control_message: IngestControlMessage, stage_config: TextSp
 
     # Remove the "split" task to obtain task-specific configuration.
     task_config = remove_task_by_type(control_message, "split")
-    logger.debug("Extracted task config: %s", task_config)
+    logger.debug("Extracted task config: %s", sanitize_for_logging(task_config))
 
     # Transform the DataFrame (split text and tokenize).
     df_updated = transform_text_split_and_tokenize_internal(
@@ -110,10 +117,10 @@ def text_splitter_fn(control_message: IngestControlMessage, stage_config: TextSp
         transform_config=stage_config,
         execution_trace_log=None,
     )
-    logger.info("TextSplitterStage.on_data: Transformation complete. Updated payload has %d rows.", len(df_updated))
+    logger.debug("TextSplitterStage.on_data: Transformation complete. Updated payload has %d rows.", len(df_updated))
 
     # Update the message payload.
     control_message.payload(df_updated)
-    logger.info("TextSplitterStage.on_data: Finished processing, returning updated message.")
+    logger.debug("TextSplitterStage.on_data: Finished processing, returning updated message.")
 
     return control_message

nv_ingest/framework/orchestration/ray/stages/utility/throughput_monitor.py

@@ -4,11 +4,16 @@
 
 import logging
 import time
-from typing import Any
+from typing import Any, Optional
 from pydantic import BaseModel
 import ray
 
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control.udf_intercept import udf_intercept_hook
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -22,12 +27,15 @@ class ThroughputMonitorStage(RayActorStage):
     It also adds the throughput as metadata on the control message before passing it on.
     """
 
-    def __init__(self, config: BaseModel) -> None:
+    def __init__(self, config: BaseModel, stage_name: Optional[str] = None) -> None:
         # Initialize base attributes (e.g., self._running, self.start_time) via the base class.
-        super().__init__(config)
+        super().__init__(config, stage_name=stage_name)
         self.count = 0
         self.last_emit_time = None  # Timestamp when the last throughput measure was emitted
 
+    @nv_ingest_node_failure_try_except()
+    @traceable()
+    @udf_intercept_hook()
     async def on_data(self, message: Any) -> Any:
         """
         Process an incoming control message. Increment the internal counter and, every 100 messages,

nv_ingest/framework/orchestration/ray/util/pipeline/pid_controller.py

@@ -10,7 +10,6 @@ from typing import Dict, Any, List, Tuple, Optional
 
 from nv_ingest_api.util.system.hardware_info import SystemResourceProbe
 
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # --- Constants ---
@@ -259,7 +258,7 @@ class ResourceConstraintManager:
         else:
             self.core_based_replica_limit = None  # Treat as unlimited if detection failed
 
-        logger.info(
+        logger.debug(
             f"[ConstraintMgr] Initialized. MaxReplicas={max_replicas}, "
             f"EffectiveCoreLimit={self.available_cores:.2f} "  # Log the potentially fractional value
             f"(Method: {self.core_detection_details.get('detection_method')}), "