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

nv_ingest/pipeline/config/loaders.py

@@ -0,0 +1,229 @@
+# 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_libmode_pipeline_impl import DEFAULT_LIBMODE_PIPELINE_YAML
+from nv_ingest.pipeline.default_pipeline_impl import DEFAULT_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_pipeline_config() -> PipelineConfigSchema:
+    """
+    Load and validate the embedded default (non-libmode) pipeline configuration.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        Validated default pipeline configuration.
+
+    Raises
+    ------
+    ValueError
+        If the default YAML cannot be parsed or validated.
+    """
+    logger.info("Loading embedded default pipeline configuration")
+
+    substituted_content = substitute_env_vars_in_yaml_content(DEFAULT_PIPELINE_YAML)
+
+    try:
+        processed_config = yaml.safe_load(substituted_content)
+    except yaml.YAMLError as e:
+        error_message = (
+            f"Failed to parse embedded default pipeline YAML after environment variable substitution. Error: {e}"
+        )
+        raise ValueError(error_message) from e
+
+    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:
+        # For non-libmode, fall back to embedded default pipeline implementation
+        return load_default_pipeline_config()

nv_ingest/pipeline/config/replica_resolver.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Runtime replica resolution for static scaling mode.
+
+This module provides functionality to resolve replica counts for stages using
+non-static strategies when dynamic scaling is disabled, ensuring total memory
+consumption stays within the static_memory_threshold.
+"""
+
+import logging
+import os
+from typing import List
+from copy import deepcopy
+
+from nv_ingest.pipeline.pipeline_schema import (
+    PipelineConfigSchema,
+    StageConfig,
+    ReplicaCalculationStrategy,
+    ReplicaStrategyConfig,
+)
+from nv_ingest_api.util.system.hardware_info import SystemResourceProbe
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_static_replicas(pipeline_config: PipelineConfigSchema) -> PipelineConfigSchema:
+    """
+    Resolve static replica counts for all stages when dynamic scaling is disabled.
+
+    This function calculates the static replica counts for stages using non-static
+    strategies, ensuring the total memory consumption stays within the configured
+    static_memory_threshold. If the total exceeds the threshold, all non-static
+    stages are scaled down proportionally (minimum 1 replica each).
+
+    Parameters
+    ----------
+    pipeline_config : PipelineConfigSchema
+        The pipeline configuration with potentially unresolved replica strategies.
+
+    Returns
+    -------
+    PipelineConfigSchema
+        A new pipeline configuration with all static replica counts resolved.
+    """
+    # Only resolve if dynamic scaling is disabled
+    if not pipeline_config.pipeline.disable_dynamic_scaling:
+        logger.debug("Dynamic scaling enabled, skipping static replica resolution")
+        return pipeline_config
+
+    logger.info("Resolving static replica counts for disabled dynamic scaling mode")
+
+    # Create a deep copy to avoid modifying the original config
+    resolved_config = deepcopy(pipeline_config)
+
+    # Get system resource information
+    system_probe = SystemResourceProbe()
+    total_memory_mb = system_probe.total_memory_mb
+    available_memory_mb = int(total_memory_mb * resolved_config.pipeline.static_memory_threshold)
+
+    logger.info(
+        f"System memory: {total_memory_mb}MB, available for static replicas: {available_memory_mb}MB "
+        f"(threshold: {resolved_config.pipeline.static_memory_threshold:.1%})"
+    )
+
+    # Find stages with non-static strategies and calculate their baseline replica counts
+    non_static_stages = []
+    total_memory_demand_mb = 0
+
+    for stage in resolved_config.stages:
+        if stage.replicas and stage.replicas.static_replicas:
+            if isinstance(stage.replicas.static_replicas, ReplicaStrategyConfig):
+                strategy_config = stage.replicas.static_replicas
+                baseline_replicas = _calculate_baseline_static_replicas(
+                    stage, strategy_config, system_probe, resolved_config.pipeline.static_memory_threshold
+                )
+
+                memory_per_replica_mb = strategy_config.memory_per_replica_mb or 0
+                stage_memory_demand = baseline_replicas * memory_per_replica_mb
+
+                non_static_stages.append(
+                    {
+                        "stage": stage,
+                        "strategy_config": strategy_config,
+                        "baseline_replicas": baseline_replicas,
+                        "memory_per_replica_mb": memory_per_replica_mb,
+                        "baseline_memory_demand_mb": stage_memory_demand,
+                    }
+                )
+
+                total_memory_demand_mb += stage_memory_demand
+
+                logger.debug(
+                    f"Stage '{stage.name}': {baseline_replicas} replicas × "
+                    f"{memory_per_replica_mb}MB = {stage_memory_demand}MB"
+                )
+
+    if not non_static_stages:
+        logger.info("No stages with non-static strategies found")
+        return resolved_config
+
+    logger.info(f"Total baseline memory demand: {total_memory_demand_mb}MB from {len(non_static_stages)} stages")
+
+    # Optional bypass of global memory-based scale down via environment variable
+    bypass_env = os.getenv("NV_INGEST_BYPASS_STATIC_MEMORY_SCALE_DOWN", "").strip().lower()
+    bypass_scale_down = bypass_env in ("1", "true", "yes", "on")
+
+    # Check if we need to scale down (unless bypassed)
+    if bypass_scale_down:
+        logger.warning(
+            "Bypassing static memory-based replica scale-down due to NV_INGEST_BYPASS_STATIC_MEMORY_SCALE_DOWN"
+        )
+        scaling_factor = 1.0
+    elif total_memory_demand_mb <= available_memory_mb:
+        logger.info("Memory demand within threshold, applying baseline replica counts")
+        scaling_factor = 1.0
+    else:
+        # Calculate scaling factor to fit within memory threshold
+        scaling_factor = available_memory_mb / total_memory_demand_mb
+        logger.warning(
+            f"Memory demand exceeds threshold by {((total_memory_demand_mb / available_memory_mb) - 1) * 100:.1f}%, "
+            f"scaling down by factor of {scaling_factor:.3f}"
+        )
+
+    # Apply the resolved replica counts
+    total_actual_memory_mb = 0
+    for stage_info in non_static_stages:
+        stage = stage_info["stage"]
+        baseline_replicas = stage_info["baseline_replicas"]
+        memory_per_replica_mb = stage_info["memory_per_replica_mb"]
+
+        # Calculate scaled replica count (minimum 1)
+        scaled_replicas = max(1, int(baseline_replicas * scaling_factor))
+        actual_memory_mb = scaled_replicas * memory_per_replica_mb
+        total_actual_memory_mb += actual_memory_mb
+
+        # Replace the strategy config with a static replica count
+        stage.replicas.static_replicas = scaled_replicas
+
+        logger.info(
+            f"Stage '{stage.name}': {baseline_replicas} → {scaled_replicas} replicas " f"({actual_memory_mb}MB)"
+        )
+
+    logger.info(
+        f"Total actual memory allocation: {total_actual_memory_mb}MB "
+        f"({(total_actual_memory_mb / total_memory_mb) * 100:.1f}% of system memory)"
+    )
+
+    return resolved_config
+
+
+def _calculate_baseline_static_replicas(
+    stage: StageConfig,
+    strategy_config: ReplicaStrategyConfig,
+    system_probe: SystemResourceProbe,
+    static_memory_threshold: float = 0.75,
+) -> int:
+    """
+    Calculate the baseline static replica count for a stage based on its strategy.
+
+    Parameters
+    ----------
+    stage : StageConfig
+        The stage configuration.
+    strategy_config : ReplicaStrategyConfig
+        The replica strategy configuration.
+    system_probe : SystemResourceProbe
+        System resource information.
+    static_memory_threshold : float, optional
+        The global static memory threshold (default: 0.75).
+
+    Returns
+    -------
+    int
+        The calculated baseline replica count.
+    """
+    strategy = strategy_config.strategy
+
+    if strategy == ReplicaCalculationStrategy.STATIC:
+        return strategy_config.value or 1
+
+    elif strategy == ReplicaCalculationStrategy.CPU_PERCENTAGE:
+        cpu_percent = strategy_config.cpu_percent or 0.5
+        limit = strategy_config.limit or system_probe.cpu_count
+        calculated = max(1, int(system_probe.cpu_count * cpu_percent))
+        return min(calculated, limit)
+
+    elif strategy == ReplicaCalculationStrategy.MEMORY_THRESHOLDING:
+        # For memory thresholding, use a conservative approach for static mode
+        memory_per_replica_mb = strategy_config.memory_per_replica_mb or 1000
+        available_memory_mb = int(system_probe.total_memory_mb * 0.7)  # Conservative 70%
+        calculated = max(1, available_memory_mb // memory_per_replica_mb)
+        limit = strategy_config.limit or calculated
+        return min(calculated, limit)
+
+    elif strategy == ReplicaCalculationStrategy.MEMORY_STATIC_GLOBAL_PERCENT:
+        # Use the global static memory threshold for calculation
+        memory_per_replica_mb = strategy_config.memory_per_replica_mb or 1000
+        available_memory_mb = int(system_probe.total_memory_mb * static_memory_threshold)
+        calculated = max(1, available_memory_mb // memory_per_replica_mb)
+        limit = strategy_config.limit or calculated
+        return min(calculated, limit)
+
+    else:
+        logger.warning(f"Unknown replica strategy '{strategy}' for stage '{stage.name}', defaulting to 1 replica")
+        return 1
+
+
+def get_memory_intensive_stages(pipeline_config: PipelineConfigSchema) -> List[str]:
+    """
+    Identify stages that are memory-intensive and may need special handling.
+
+    Parameters
+    ----------
+    pipeline_config : PipelineConfigSchema
+        The pipeline configuration.
+
+    Returns
+    -------
+    List[str]
+        List of stage names that are memory-intensive.
+    """
+    memory_intensive_stages = []
+
+    for stage in pipeline_config.stages:
+        if stage.replicas and stage.replicas.static_replicas:
+            if isinstance(stage.replicas.static_replicas, ReplicaStrategyConfig):
+                strategy_config = stage.replicas.static_replicas
+                memory_per_replica_mb = strategy_config.memory_per_replica_mb or 0
+
+                # Consider stages using >5GB per replica as memory-intensive
+                if memory_per_replica_mb > 5000:
+                    memory_intensive_stages.append(stage.name)
+
+    return memory_intensive_stages