nv-ingest 2025.5.21.dev20250521__py3-none-any.whl

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

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Any
+from pydantic import BaseModel
+import ray
+
+# Import the base class for our stages.
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.schemas.framework_job_counter_schema import JobCounterSchema
+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,
+)
+
+# Import the JobCounter schema and global stats singleton.
+
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class JobCounterStage(RayActorStage):
+    """
+    A Ray actor stage that counts jobs and updates global statistics.
+
+    Based on the configuration (a JobCounterSchema instance), it increments a specific
+    statistic each time it processes a message.
+    """
+
+    def __init__(self, config: BaseModel) -> None:
+        # Ensure base attributes (e.g. self._running) are initialized.
+        super().__init__(config)
+        # 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)
+    async def on_data(self, message: Any) -> Any:
+        """
+        Process an incoming IngestControlMessage by counting jobs.
+
+        If the validated configuration name is "completed_jobs", then if the message metadata
+        indicates a failure (cm_failed == True), increments "failed_jobs"; otherwise, increments "completed_jobs".
+        For any other configuration name, it increments that statistic.
+
+        Returns the original message.
+        """
+        logger.debug(f"Performing job counter: {self.validated_config.name}")
+        try:
+            if self.validated_config.name == "completed_jobs":
+                if message.has_metadata("cm_failed") and message.get_metadata("cm_failed"):
+                    self.stats.increment_stat("failed_jobs")
+                else:
+                    self.stats.increment_stat("completed_jobs")
+                return message
+
+            self.stats.increment_stat(self.validated_config.name)
+            return message
+        except Exception as e:
+            new_message = f"on_data: Failed to run job counter. Original error: {str(e)}"
+            logger.exception(new_message)
+            raise type(e)(new_message) from e

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

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

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

@@ -0,0 +1,205 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from datetime import datetime
+from typing import Any, Optional
+
+import ray
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+from opentelemetry.trace import NonRecordingSpan
+from opentelemetry.trace import SpanContext
+from opentelemetry.trace import Status
+from opentelemetry.trace import StatusCode
+from opentelemetry.trace import TraceFlags
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.schemas.framework_otel_tracer_schema import OpenTelemetryTracerSchema
+from nv_ingest_api.util.exception_handlers.decorators import nv_ingest_node_failure_try_except
+
+from nv_ingest_api.internal.primitives.tracing.logging import TaskResultStatus
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage
+
+
+@ray.remote
+class OpenTelemetryTracerStage(RayActorStage):
+    """
+    A Ray actor stage that collects and exports traces to OpenTelemetry.
+
+    This stage uses OpenTelemetry to trace the execution of tasks within the system.
+    It creates spans for tasks and exports them to a configured OpenTelemetry endpoint.
+    """
+
+    def __init__(self, config: OpenTelemetryTracerSchema) -> None:
+        super().__init__(config)
+
+        # self._logger.info(f"[Telemetry] Initializing OpenTelemetry tracer stage with config: {config}")
+
+        self.validated_config: OpenTelemetryTracerSchema = config
+        self.resource = Resource(attributes={"service.name": "nv-ingest"})
+        self.otlp_exporter = OTLPSpanExporter(endpoint=self.validated_config.otel_endpoint, insecure=True)
+        self.span_processor = BatchSpanProcessor(self.otlp_exporter)
+
+        trace.set_tracer_provider(TracerProvider(resource=self.resource))
+        trace.get_tracer_provider().add_span_processor(self.span_processor)
+
+        self.tracer = trace.get_tracer(__name__)
+
+    def collect_timestamps(self, message):
+        job_id = message.get_metadata("job_id")
+        if isinstance(job_id, str) and len(job_id) == 36:
+            trace_id = uuid_to_trace_id(job_id)
+        elif isinstance(job_id, str):
+            trace_id = int(job_id, 16)
+        else:
+            trace_id = RandomIdGenerator().generate_trace_id()
+
+        span_id = RandomIdGenerator().generate_span_id()
+        timestamps = extract_timestamps_from_message(message)
+
+        flattened = [x for t in timestamps.values() for x in t]
+        if not flattened:
+            self._logger.debug("No timestamps found for message; skipping tracing.")
+            return
+
+        start_time = min(flattened)
+        end_time = max(flattened)
+
+        self._logger.debug(f"[Telemetry] trace_id: {trace_id}, span_id: {span_id}")
+
+        span_context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=True,
+            trace_flags=TraceFlags(0x01),
+        )
+        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)
+
+        if message.has_metadata("cm_failed") and message.get_metadata("cm_failed"):
+            parent_span.set_status(Status(StatusCode.ERROR))
+        else:
+            parent_span.set_status(Status(StatusCode.OK))
+
+        try:
+            parent_span.add_event("start", timestamp=start_time)
+            parent_span.add_event("end", timestamp=end_time)
+        finally:
+            parent_span.end(end_time=end_time)
+
+        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)
+    def on_data(self, control_message: IngestControlMessage) -> Optional[Any]:
+        try:
+            do_trace_tagging = bool(control_message.get_metadata("config::add_trace_tagging"))
+
+            if not do_trace_tagging:
+                self._logger.debug("Skipping OpenTelemetry tracing, do_trace_tagging is False.")
+                return control_message
+
+            self._logger.debug("Sending telemetry data to OpenTelemetry")
+
+            self.collect_timestamps(control_message)
+
+            return control_message
+        except Exception as e:
+            self._logger.warning(f"Error in OpenTelemetry tracer stage: {e}")
+            raise e
+
+
+def extract_timestamps_from_message(message):
+    timestamps = {}
+    dedup_counter = {}
+
+    for key, val in message.filter_timestamp("trace::exit::").items():
+        exit_key = key
+        entry_key = exit_key.replace("trace::exit::", "trace::entry::")
+
+        task_name = key.replace("trace::exit::", "")
+        if task_name in dedup_counter:
+            dedup_counter[task_name] += 1
+            task_name = task_name + "_" + str(dedup_counter[task_name])
+        else:
+            dedup_counter[task_name] = 0
+
+        ts_entry = message.get_timestamp(entry_key)
+        if ts_entry is None:
+            continue
+
+        ts_exit = (
+            message.get_timestamp(exit_key) or datetime.now()
+        )  # When a job fails, it may not have exit time. Default to current time.
+        ts_entry_ns = int(ts_entry.timestamp() * 1e9)
+        ts_exit_ns = int(ts_exit.timestamp() * 1e9)
+
+        timestamps[task_name] = (ts_entry_ns, ts_exit_ns)
+
+    return timestamps
+
+
+def extract_annotated_task_results(message):
+    task_results = {}
+    for key in message.list_metadata():
+        if not key.startswith("annotation::"):
+            continue
+        task = message.get_metadata(key)
+        if not (("task_id" in task) and ("task_result" in task)):
+            continue
+        task_id = task["task_id"]
+        task_result = task["task_result"]
+        task_results[task_id] = task_result
+
+    return task_results
+
+
+def create_span_with_timestamps(tracer, parent_span, message) -> int:
+    timestamps = extract_timestamps_from_message(message)
+    task_results = extract_annotated_task_results(message)
+
+    ctx_store = {}
+    event_counter = 0
+    child_ctx = trace.set_span_in_context(parent_span)
+
+    for task_name, (ts_entry, ts_exit) in sorted(timestamps.items(), key=lambda x: x[1]):
+        main_task, *subtask = task_name.split("::", 1)
+        subtask = "::".join(subtask)
+
+        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)
+
+        span.add_event("entry", timestamp=ts_entry)
+        span.add_event("exit", timestamp=ts_exit)
+        event_counter += 2
+
+        if task_name in task_results:
+            task_result = task_results[main_task]
+            if task_result == TaskResultStatus.SUCCESS.value:
+                span.set_status(Status(StatusCode.OK))
+            if task_result == TaskResultStatus.FAILURE.value:
+                span.set_status(Status(StatusCode.ERROR))
+
+        ctx_store[task_name] = (span, ts_exit)
+
+    for _, (span, ts_exit) in ctx_store.items():
+        span.end(end_time=ts_exit)
+
+    return event_counter
+
+
+def uuid_to_trace_id(uuid_str: str) -> int:
+    """Convert a UUID-like string to an integer OpenTelemetry trace ID."""
+    if not isinstance(uuid_str, str) or len(uuid_str) != 36:
+        raise ValueError("UUID must be a 36-character string with hyphens")
+
+    return int(uuid_str.replace("-", ""), 16)

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

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

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

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import pprint
+from typing import Any
+
+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
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.transform.transform_image_caption_schema import ImageCaptionExtractionSchema
+from nv_ingest_api.internal.transform.caption_image import transform_image_create_vlm_caption_internal
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class ImageCaptionTransformStage(RayActorStage):
+    """
+    A Ray actor stage that extracts image captions from a DataFrame payload.
+
+    This stage validates its configuration (using ImageCaptionExtractionSchema), then processes the DataFrame
+    via transform_image_create_vlm_caption_internal. The updated DataFrame and any extraction trace info
+    are stored in the control message.
+    """
+
+    def __init__(self, config: ImageCaptionExtractionSchema) -> None:
+        super().__init__(config)
+        try:
+            self.validated_config = config
+            logger.info("ImageCaptionTransformStage configuration validated.")
+        except Exception as e:
+            logger.exception("Error validating caption extraction config")
+            raise e
+
+    @traceable("image_captioning")
+    @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.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The incoming message containing the DataFrame payload.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with the extracted captions.
+        """
+        logger.info("ImageCaptionTransformStage.on_data: Starting image caption extraction.")
+
+        # Retrieve the DataFrame payload.
+        df_payload = control_message.payload()
+        logger.debug("ImageCaptionTransformStage: Payload extracted with %d rows.", len(df_payload))
+
+        # 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))
+
+        # Call the caption extraction function.
+        new_df = transform_image_create_vlm_caption_internal(
+            df_payload, task_config=task_config, transform_config=self.validated_config
+        )
+        logger.info("Image caption extraction completed. New payload has %d rows.", len(new_df))
+
+        # Update the control message with the new DataFrame.
+        control_message.payload(new_df)
+        # Optionally, annotate the control message with extraction trace info.
+        # control_message.set_metadata("caption_extraction_trace", execution_trace_log)
+        logger.info("ImageCaptionTransformStage.on_data: Updated control message and returning.")
+        return control_message

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

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import pprint
+from typing import Any
+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
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.transform.transform_text_embedding_schema import TextEmbeddingSchema
+from nv_ingest_api.internal.transform.embed_text import transform_create_text_embeddings_internal
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class TextEmbeddingTransformStage(RayActorStage):
+    """
+    A Ray actor stage that extracts text embeddings from a DataFrame payload.
+
+    This stage uses the validated configuration (TextEmbeddingSchema) to process the DataFrame
+    and generate text embeddings. The resulting DataFrame is set back on the message, and any
+    trace or extraction metadata is added.
+    """
+
+    def __init__(self, config: TextEmbeddingSchema) -> None:
+        super().__init__(config)
+        try:
+            self.validated_config = config
+            logger.info("TextEmbeddingTransformStage configuration validated successfully.")
+        except Exception as e:
+            logger.exception("Error validating text embedding extractor config")
+            raise e
+
+    @traceable("text_embedding")
+    @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:
+        """
+        Process the control message by generating text embeddings.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The incoming message containing the DataFrame payload.
+
+        Returns
+        -------
+        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))
+
+        # 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))
+
+        # 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

@@ -0,0 +1,74 @@
+import logging
+from typing import Any
+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
+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,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class TextSplitterStage(RayActorStage):
+    """
+    A Ray actor stage that splits documents into smaller parts based on specified criteria.
+
+    This stage extracts the DataFrame payload from an IngestControlMessage, removes the "split"
+    task (if present) to obtain the task configuration, and then calls the internal text splitting
+    and tokenization logic. The updated DataFrame is then set back into the message.
+    """
+
+    def __init__(self, config: TextSplitterSchema) -> None:
+        super().__init__(config)
+        # Store the validated configuration (assumed to be an instance of TextSplitterSchema)
+        self.validated_config: TextSplitterSchema = config
+        logger.info("TextSplitterStage initialized with config: %s", config)
+
+    @traceable("text_splitter")
+    @filter_by_task(["split"])
+    @nv_ingest_node_failure_try_except(annotation_id="text_splitter", raise_on_failure=False)
+    def on_data(self, message: Any) -> Any:
+        """
+        Process an incoming IngestControlMessage by splitting and tokenizing its text.
+
+        Parameters
+        ----------
+        message : IngestControlMessage
+            The incoming message containing the payload DataFrame.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with its payload transformed.
+        """
+
+        # Extract the DataFrame payload.
+        df_payload = message.payload()
+        logger.debug("Extracted payload with %d rows.", len(df_payload))
+
+        # 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)
+
+        # Transform the DataFrame (split text and tokenize).
+        df_updated = transform_text_split_and_tokenize_internal(
+            df_transform_ledger=df_payload,
+            task_config=task_config,
+            transform_config=self.validated_config,
+            execution_trace_log=None,
+        )
+        logger.info("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.")
+
+        return message

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

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

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

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import time
+from typing import Any
+from pydantic import BaseModel
+import ray
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class ThroughputMonitorStage(RayActorStage):
+    """
+    A Ray actor stage that monitors throughput by counting messages.
+
+    Every 100 messages, it calculates the throughput (messages per second) and logs the measure.
+    It also adds the throughput as metadata on the control message before passing it on.
+    """
+
+    def __init__(self, config: BaseModel) -> None:
+        # Initialize base attributes (e.g., self._running, self.start_time) via the base class.
+        super().__init__(config)
+        self.count = 0
+        self.last_emit_time = None  # Timestamp when the last throughput measure was emitted
+
+    async def on_data(self, message: Any) -> Any:
+        """
+        Process an incoming control message. Increment the internal counter and, every 100 messages,
+        calculate and log the throughput. The throughput value is also added to the message metadata.
+
+        Parameters
+        ----------
+        message : Any
+            The incoming control message.
+
+        Returns
+        -------
+        Any
+            The (possibly modified) control message.
+        """
+        self.count += 1
+        if self.last_emit_time is None:
+            self.last_emit_time = time.time()
+
+        if self.count % 1000 == 0:
+            now = time.time()
+            elapsed = now - self.last_emit_time
+            throughput = 1000 / elapsed if elapsed > 0 else 0
+            logger.warning(
+                f"ThroughputMonitorStage: Processed {self.count} messages. Throughput: {throughput:.2f} messages/sec"
+            )
+            try:
+                # Attempt to add throughput information to the message metadata.
+                message.set_metadata("throughput", throughput)
+            except Exception:
+                # If the message doesn't support metadata, skip.
+                pass
+            self.last_emit_time = now
+
+        return message

nv_ingest/framework/orchestration/ray/util/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

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

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0