nv-ingest 2025.8.14.dev20250814__py3-none-any.whl → 2025.8.16.dev20250816__py3-none-any.whl

nv_ingest/framework/util/flow_control/udf_intercept.py

@@ -0,0 +1,352 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import hashlib
+import inspect
+import logging
+import os
+import time
+from typing import Dict, List, Any, Optional, Callable
+from dataclasses import dataclass
+
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_all_tasks_by_type
+from nv_ingest_api.util.imports.callable_signatures import ingest_callable_signature
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CachedUDF:
+    """Cached UDF function with metadata"""
+
+    function: callable
+    function_name: str
+    signature_validated: bool
+    created_at: float
+    last_used: float
+    use_count: int
+
+
+class UDFCache:
+    """LRU cache for compiled and validated UDF functions"""
+
+    def __init__(self, max_size: int = 128, ttl_seconds: Optional[int] = 3600):
+        self.max_size = max_size
+        self.ttl_seconds = ttl_seconds
+        self.cache: Dict[str, CachedUDF] = {}
+        self.access_order: List[str] = []  # For LRU tracking
+
+    def _generate_cache_key(self, udf_function_str: str, udf_function_name: str) -> str:
+        """Generate cache key from UDF string and function name"""
+        content = f"{udf_function_str.strip()}:{udf_function_name}"
+        return hashlib.sha256(content.encode()).hexdigest()
+
+    def _evict_lru(self):
+        """Remove least recently used item"""
+        if self.access_order:
+            lru_key = self.access_order.pop(0)
+            self.cache.pop(lru_key, None)
+
+    def _cleanup_expired(self):
+        """Remove expired entries if TTL is configured"""
+        if not self.ttl_seconds:
+            return
+
+        current_time = time.time()
+        expired_keys = [
+            key for key, cached_udf in self.cache.items() if current_time - cached_udf.created_at > self.ttl_seconds
+        ]
+
+        for key in expired_keys:
+            self.cache.pop(key, None)
+            if key in self.access_order:
+                self.access_order.remove(key)
+
+    def get(self, udf_function_str: str, udf_function_name: str) -> Optional[CachedUDF]:
+        """Get cached UDF function if available"""
+        self._cleanup_expired()
+
+        cache_key = self._generate_cache_key(udf_function_str, udf_function_name)
+
+        if cache_key in self.cache:
+            # Update access tracking
+            if cache_key in self.access_order:
+                self.access_order.remove(cache_key)
+            self.access_order.append(cache_key)
+
+            # Update usage stats
+            cached_udf = self.cache[cache_key]
+            cached_udf.last_used = time.time()
+            cached_udf.use_count += 1
+
+            return cached_udf
+
+        return None
+
+    def put(
+        self, udf_function_str: str, udf_function_name: str, function: callable, signature_validated: bool = True
+    ) -> str:
+        """Cache a compiled and validated UDF function"""
+        cache_key = self._generate_cache_key(udf_function_str, udf_function_name)
+
+        # Evict LRU if at capacity
+        while len(self.cache) >= self.max_size:
+            self._evict_lru()
+
+        current_time = time.time()
+        cached_udf = CachedUDF(
+            function=function,
+            function_name=udf_function_name,
+            signature_validated=signature_validated,
+            created_at=current_time,
+            last_used=current_time,
+            use_count=1,
+        )
+
+        self.cache[cache_key] = cached_udf
+        self.access_order.append(cache_key)
+
+        return cache_key
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get cache statistics"""
+        total_uses = sum(udf.use_count for udf in self.cache.values())
+        most_used = max(self.cache.values(), key=lambda x: x.use_count, default=None)
+        return {
+            "size": len(self.cache),
+            "max_size": self.max_size,
+            "total_uses": total_uses,
+            "most_used_function": most_used.function_name if most_used else None,
+            "most_used_count": most_used.use_count if most_used else 0,
+        }
+
+
+# Global cache instance
+_udf_cache = UDFCache(max_size=128, ttl_seconds=3600)
+
+
+def compile_and_validate_udf(udf_function_str: str, udf_function_name: str, task_num: int) -> callable:
+    """Compile and validate UDF function (extracted for caching)"""
+    # Execute the UDF function string in a controlled namespace
+    namespace: Dict[str, Any] = {}
+    try:
+        exec(udf_function_str, namespace)
+    except Exception as e:
+        raise ValueError(f"UDF task {task_num} failed to execute: {str(e)}")
+
+    # Extract the specified function from the namespace
+    if udf_function_name in namespace and callable(namespace[udf_function_name]):
+        udf_function = namespace[udf_function_name]
+    else:
+        raise ValueError(f"UDF task {task_num}: Specified UDF function '{udf_function_name}' not found or not callable")
+
+    # Validate the UDF function signature
+    try:
+        ingest_callable_signature(inspect.signature(udf_function))
+    except Exception as e:
+        raise ValueError(f"UDF task {task_num} has invalid function signature: {str(e)}")
+
+    return udf_function
+
+
+def execute_targeted_udfs(
+    control_message: IngestControlMessage, stage_name: str, directive: str
+) -> IngestControlMessage:
+    """Execute UDFs that target this stage with the given directive."""
+    # Early exit if no UDF tasks exist - check by task type, not task ID
+    udf_tasks_exist = any(task.type == "udf" for task in control_message.get_tasks())
+    if not udf_tasks_exist:
+        return control_message
+
+    # Remove all UDF tasks and get them - handle case where no tasks found
+    try:
+        all_udf_tasks = remove_all_tasks_by_type(control_message, "udf")
+    except ValueError:
+        # No UDF tasks found - this can happen due to race conditions
+        logger.debug(f"No UDF tasks found for stage '{stage_name}' directive '{directive}'")
+        return control_message
+
+    # Execute applicable UDFs and collect remaining ones
+    remaining_tasks = []
+
+    for task_properties in all_udf_tasks:
+        # Check if this UDF targets this stage with the specified directive
+        target_stage = task_properties.get("target_stage", "")
+        run_before = task_properties.get("run_before", False)
+        run_after = task_properties.get("run_after", False)
+
+        # Determine if this UDF should execute
+        should_execute = False
+        if directive == "run_before" and run_before and target_stage == stage_name:
+            should_execute = True
+        elif directive == "run_after" and run_after and target_stage == stage_name:
+            should_execute = True
+
+        if should_execute:
+            try:
+                # Get UDF function details
+                udf_function_str = task_properties.get("udf_function", "").strip()
+                udf_function_name = task_properties.get("udf_function_name", "").strip()
+                task_id = task_properties.get("task_id", "unknown")
+
+                # Skip empty UDF functions
+                if not udf_function_str:
+                    logger.debug(f"UDF task {task_id} has empty function, skipping")
+                    remaining_tasks.append(task_properties)
+                    continue
+
+                # Validate function name
+                if not udf_function_name:
+                    raise ValueError(f"UDF task {task_id} missing required 'udf_function_name' property")
+
+                # Get or compile UDF function
+                cached_udf = _udf_cache.get(udf_function_str, udf_function_name)
+                if cached_udf:
+                    udf_function = cached_udf.function
+                    logger.debug(f"UDF task {task_id}: Using cached function '{udf_function_name}'")
+                else:
+                    udf_function = compile_and_validate_udf(udf_function_str, udf_function_name, task_id)
+                    _udf_cache.put(udf_function_str, udf_function_name, udf_function)
+                    logger.debug(f"UDF task {task_id}: Cached function '{udf_function_name}'")
+
+                # Execute the UDF
+                control_message = udf_function(control_message)
+
+                # Validate return type
+                if not isinstance(control_message, IngestControlMessage):
+                    raise ValueError(f"UDF task {task_id} must return IngestControlMessage")
+
+                logger.info(f"Executed UDF {task_id} '{udf_function_name}' {directive} stage '{stage_name}'")
+
+            except Exception as e:
+                logger.error(f"UDF {task_id} failed {directive} stage '{stage_name}': {e}")
+                # Keep failed task for next stage
+                remaining_tasks.append(task_properties)
+        else:
+            # Keep non-applicable task for next stage
+            remaining_tasks.append(task_properties)
+
+    # Re-add all remaining UDF tasks
+    for task_properties in remaining_tasks:
+        from nv_ingest_api.internal.primitives.control_message_task import ControlMessageTask
+
+        task = ControlMessageTask(type="udf", id=task_properties.get("task_id", "unknown"), properties=task_properties)
+        control_message.add_task(task)
+
+    return control_message
+
+
+def remove_task_by_id(control_message: IngestControlMessage, task_id: str) -> IngestControlMessage:
+    """Remove a specific task by ID from the control message"""
+    try:
+        control_message.remove_task(task_id)
+    except RuntimeError as e:
+        logger.warning(f"Could not remove task {task_id}: {e}")
+
+    return control_message
+
+
+def udf_intercept_hook(stage_name: Optional[str] = None, enable_run_before: bool = True, enable_run_after: bool = True):
+    """
+    Decorator that executes UDFs targeted at this stage.
+
+    This decorator integrates with the existing UDF system, providing full
+    UDF compilation, caching, and execution capabilities. UDFs can target
+    specific stages using run_before or run_after directives.
+
+    Args:
+        stage_name: Name of the stage (e.g., "image_dedup", "text_extract").
+                   If None, will attempt to use self.stage_name from the decorated method's instance.
+        enable_run_before: Whether to execute UDFs with run_before=True (default: True)
+        enable_run_after: Whether to execute UDFs with run_after=True (default: True)
+
+    Examples:
+        # Automatic stage name detection (recommended)
+        @traceable("image_deduplication")
+        @udf_intercept_hook()  # Uses self.stage_name automatically
+        @filter_by_task(required_tasks=["dedup"])
+        def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+            return control_message
+
+        # Explicit stage name (fallback/override)
+        @traceable("data_sink")
+        @udf_intercept_hook("data_sink", enable_run_after=False)
+        @filter_by_task(required_tasks=["store"])
+        def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+            return control_message
+
+        # Only run_after UDFs (e.g., for source stages)
+        @traceable("data_source")
+        @udf_intercept_hook(enable_run_before=False)  # Uses self.stage_name automatically
+        def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+            return control_message
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Check if UDF processing is globally disabled
+            if os.getenv("INGEST_DISABLE_UDF_PROCESSING"):
+                logger.debug("UDF processing is disabled via INGEST_DISABLE_UDF_PROCESSING environment variable")
+                return func(*args, **kwargs)
+
+            # Determine the stage name to use
+            resolved_stage_name = stage_name
+
+            # If no explicit stage_name provided, try to get it from self.stage_name
+            if resolved_stage_name is None and len(args) >= 1:
+                stage_instance = args[0]  # 'self' in method calls
+                if hasattr(stage_instance, "stage_name") and stage_instance.stage_name:
+                    resolved_stage_name = stage_instance.stage_name
+                    logger.debug(f"Using auto-detected stage name: '{resolved_stage_name}'")
+                else:
+                    logger.warning(
+                        "No stage_name provided and could not auto-detect from instance. Skipping UDF intercept."
+                    )
+                    return func(*args, **kwargs)
+            elif resolved_stage_name is None:
+                logger.warning(
+                    "No stage_name provided and no instance available for auto-detection. Skipping UDF intercept."
+                )
+                return func(*args, **kwargs)
+
+            # Extract control_message from args (handle both self.method and function cases)
+            control_message = None
+            if len(args) >= 2 and hasattr(args[1], "get_tasks"):
+                control_message = args[1]  # self.method case
+                args_list = list(args)
+            elif len(args) >= 1 and hasattr(args[0], "get_tasks"):
+                control_message = args[0]  # function case
+                args_list = list(args)
+
+            if control_message:
+                # Execute UDFs that should run before this stage (if enabled)
+                if enable_run_before:
+                    control_message = execute_targeted_udfs(control_message, resolved_stage_name, "run_before")
+                    # Update args with modified control_message
+                    if len(args) >= 2 and hasattr(args[1], "get_tasks"):
+                        args_list[1] = control_message
+                    else:
+                        args_list[0] = control_message
+
+                # Execute the original stage logic
+                result = func(*tuple(args_list), **kwargs)
+
+                # Execute UDFs that should run after this stage (if enabled)
+                if enable_run_after and hasattr(result, "get_tasks"):  # Result is control_message
+                    result = execute_targeted_udfs(result, resolved_stage_name, "run_after")
+
+                return result
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def get_udf_cache_stats() -> Dict[str, Any]:
+    """Get UDF cache performance statistics"""
+    return _udf_cache.get_stats()

nv_ingest/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

nv_ingest/pipeline/config/__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/pipeline/config/loaders.py

@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Configuration loading and management functions for pipeline execution.
+
+This module provides declarative functions for loading, validating, and applying
+runtime overrides to pipeline configurations, replacing imperative inline logic.
+"""
+
+import logging
+import yaml
+from typing import Optional
+
+from nv_ingest.pipeline.pipeline_schema import PipelineConfigSchema
+from nv_ingest.pipeline.default_pipeline_impl import DEFAULT_LIBMODE_PIPELINE_YAML
+from nv_ingest.framework.orchestration.execution.options import PipelineRuntimeOverrides
+from nv_ingest_api.util.string_processing.yaml import substitute_env_vars_in_yaml_content
+
+logger = logging.getLogger(__name__)
+
+
+def load_pipeline_config(config_path: str) -> PipelineConfigSchema:
+    """
+    Load a pipeline configuration file, substituting environment variables.
+
+    Parameters
+    ----------
+    config_path : str
+        The path to the YAML configuration file.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        A validated PipelineConfigSchema object.
+
+    Raises
+    ------
+    ValueError
+        If the YAML file cannot be parsed after environment variable substitution.
+    """
+    logger.info(f"Loading pipeline configuration from: {config_path}")
+
+    # Read the raw YAML file content
+    with open(config_path, "r") as f:
+        raw_content = f.read()
+
+    # Substitute all environment variable placeholders using the utility function
+    substituted_content = substitute_env_vars_in_yaml_content(raw_content)
+
+    # Parse the substituted content with PyYAML, with error handling
+    try:
+        processed_config = yaml.safe_load(substituted_content)
+    except yaml.YAMLError as e:
+        error_message = (
+            f"Failed to parse YAML after environment variable substitution. "
+            f"Error: {e}\n\n"
+            f"--- Substituted Content ---\n{substituted_content}\n---------------------------"
+        )
+        raise ValueError(error_message) from e
+
+    # Pydantic validates the clean, substituted data against the schema
+    return PipelineConfigSchema(**processed_config)
+
+
+def load_default_libmode_config() -> PipelineConfigSchema:
+    """
+    Load and validate the default libmode pipeline configuration.
+
+    This function loads the embedded default libmode pipeline YAML,
+    performs environment variable substitution, and returns a validated
+    configuration object.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        Validated default libmode pipeline configuration.
+
+    Raises
+    ------
+    ValueError
+        If the default YAML cannot be parsed or validated.
+    """
+    logger.info("Loading default libmode pipeline configuration")
+
+    # Substitute environment variables in the YAML content
+    substituted_content = substitute_env_vars_in_yaml_content(DEFAULT_LIBMODE_PIPELINE_YAML)
+
+    # Parse the substituted content with PyYAML
+    try:
+        processed_config = yaml.safe_load(substituted_content)
+    except yaml.YAMLError as e:
+        error_message = (
+            f"Failed to parse default libmode pipeline YAML after environment variable substitution. " f"Error: {e}"
+        )
+        raise ValueError(error_message) from e
+
+    # Create and return validated PipelineConfigSchema
+    return PipelineConfigSchema(**processed_config)
+
+
+def apply_runtime_overrides(config: PipelineConfigSchema, overrides: PipelineRuntimeOverrides) -> PipelineConfigSchema:
+    """
+    Apply runtime parameter overrides to a pipeline configuration.
+
+    This function creates a copy of the provided configuration and applies
+    any non-None override values to the pipeline runtime settings.
+
+    Parameters
+    ----------
+    config : PipelineConfigSchema
+        Base pipeline configuration to modify.
+    overrides : PipelineRuntimeOverrides
+        Runtime overrides to apply. Only non-None values are applied.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        Modified configuration with overrides applied.
+    """
+    # Create a copy to avoid modifying the original
+    modified_config = config.model_copy(deep=True)
+
+    # Apply overrides if provided
+    if overrides.disable_dynamic_scaling is not None:
+        modified_config.pipeline.disable_dynamic_scaling = overrides.disable_dynamic_scaling
+        logger.debug(f"Applied dynamic scaling override: {overrides.disable_dynamic_scaling}")
+
+    if overrides.dynamic_memory_threshold is not None:
+        modified_config.pipeline.dynamic_memory_threshold = overrides.dynamic_memory_threshold
+        logger.debug(f"Applied memory threshold override: {overrides.dynamic_memory_threshold}")
+
+    return modified_config
+
+
+def validate_pipeline_config(config: Optional[PipelineConfigSchema]) -> PipelineConfigSchema:
+    """
+    Validate and ensure a pipeline configuration is available.
+
+    This function ensures that a valid pipeline configuration is available,
+    either from the provided config or by loading the default libmode config.
+
+    Parameters
+    ----------
+    config : Optional[PipelineConfigSchema]
+        Pipeline configuration to validate, or None to load default.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        Validated pipeline configuration.
+
+    Raises
+    ------
+    ValueError
+        If config is None and default config cannot be loaded.
+    """
+    if config is None:
+        return load_default_libmode_config()
+
+    # Config is already validated by Pydantic, just return it
+    return config
+
+
+def resolve_pipeline_config(provided_config: Optional[PipelineConfigSchema], libmode: bool) -> PipelineConfigSchema:
+    """
+    Resolve the final pipeline configuration from inputs.
+
+    This function implements the configuration resolution logic:
+    - If config provided: use it
+    - If libmode=True and no config: load default libmode config
+    - If libmode=False and no config: raise error
+
+    Parameters
+    ----------
+    provided_config : Optional[PipelineConfigSchema]
+        User-provided pipeline configuration, or None.
+    libmode : bool
+        Whether to allow loading default libmode configuration.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        Resolved and validated pipeline configuration.
+
+    Raises
+    ------
+    ValueError
+        If no config provided and libmode=False.
+    """
+    if provided_config is not None:
+        return provided_config
+
+    if libmode:
+        return load_default_libmode_config()
+    else:
+        raise ValueError("pipeline_config must be provided when libmode is False")