nemo-evaluator-launcher 0.1.28__py3-none-any.whl

nemo_evaluator_launcher/common/logging_utils.py

@@ -0,0 +1,357 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Logging configuration module for nemo-evaluator-launcher.
+
+This module provides a centralized logging configuration using structlog that outputs
+to both stderr and a log file. All modules should import and use the logger from this
+module to ensure consistent logging behavior across the application.
+
+LOGGING POLICY:
+==============
+All logging in this project MUST go through this module. This is enforced by a pre-commit
+hook that checks for violations.
+
+DO NOT:
+- import structlog directly
+- import logging directly
+- call structlog.get_logger() directly
+- call logging.getLogger() directly
+
+DO:
+- from nemo_evaluator_launcher.common.logging_utils import logger
+- from nemo_evaluator_launcher.common.logging_utils import get_logger
+
+Examples:
+    # Correct
+    from nemo_evaluator_launcher.common.logging_utils import logger
+    logger.info("User logged in", user_id="12345")
+
+    # Incorrect
+    import structlog
+    logger = structlog.get_logger()
+    logger.info("User logged in")
+"""
+
+import json
+import logging
+import logging.config
+import os
+import pathlib
+import sys
+from datetime import datetime
+from pprint import pformat
+from typing import Any, Dict
+
+import structlog
+
+# If this env var is set, it will override a more standard "LOG_LEVEL". If
+# both are unset, default would be used.
+_LOG_LEVEL_ENV_VAR = "NEMO_EVALUATOR_LOG_LEVEL"
+_DEFAULT_LOG_LEVEL = "WARNING"
+_SENSITIVE_KEY_SUBSTRINGS_NORMALIZED = {
+    # Keep minimal, broad substrings
+    # NOTE: normalized: lowercased, no spaces/_/-
+    "authorization",  # covers proxy-authorization, etc.
+    "apikey",  # covers api_key, api-key, x-api-key, nvidia_api_key, ...
+    "accesskey",  # covers access_key / access-key
+    "privatekey",
+    "token",  # covers access_token, id_token, refresh_token, *_token
+    "secret",  # covers openai_client_secret, aws_secret_access_key, *_secret
+    "password",
+    "pwd",  # common shorthand
+    "passwd",  # common variant
+}
+_ALLOWLISTED_KEYS_SUBSTRINGS = {
+    # NOTE: non-normalized (for allowlisting we want more control)
+    "_tokens",  # This likely would allow us to not redact useful stuff like `limit_tokens`, `max_new_tokens`
+}
+
+
+def _mask(val: object) -> str:
+    s = str(val)
+    if len(s) <= 10:
+        return "[REDACTED]"
+    return f"{s[:2]}…{s[-2:]}"
+
+
+def _normalize(name: object) -> str:
+    if not isinstance(name, str):
+        return ""
+    s = name.lower()
+    # drop spaces, hyphens, underscores
+    return s.replace(" ", "").replace("-", "").replace("_", "")
+
+
+def _is_sensitive_key(key: object) -> bool:
+    k_norm = _normalize(key)
+    k_non_norm = str(key)
+    return any(
+        substr in k_norm for substr in _SENSITIVE_KEY_SUBSTRINGS_NORMALIZED
+    ) and not any(substr in k_non_norm for substr in _ALLOWLISTED_KEYS_SUBSTRINGS)
+
+
+def _redact_mapping(m: dict) -> dict:
+    red = {}
+    for k, v in m.items():
+        if _is_sensitive_key(k):
+            red[k] = _mask(v)
+        elif isinstance(v, dict):
+            red[k] = _redact_mapping(v)
+        else:
+            red[k] = v
+    return red
+
+
+def redact_processor(_: Any, __: str, event_dict: Dict[str, Any]) -> Dict[str, Any]:
+    if os.getenv("LOG_DISABLE_REDACTION", "").lower() in {"1", "true", "yes"}:
+        return event_dict
+    return _redact_mapping(event_dict)
+
+
+def _ensure_log_dir() -> pathlib.Path:
+    """Ensure the log directory exists and return its path."""
+    log_dir = pathlib.Path.home() / ".nemo-evaluator" / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+    return log_dir
+
+
+def _get_env_log_level() -> str:
+    """Get log level from environment variable, translating single letters to full names.
+
+    Translates:
+    - D -> DEBUG
+    - I -> INFO
+    - W -> WARNING
+    - E -> ERROR
+    - F -> CRITICAL
+
+    Returns:
+        Uppercase log level string, defaults to WARNING if not set or invalid.
+    """
+    env_level = os.getenv(_LOG_LEVEL_ENV_VAR, os.getenv("LOG_LEVEL"))
+    # If empty or unset, default
+    if not env_level:
+        env_level = _DEFAULT_LOG_LEVEL
+    env_level = env_level.upper()
+
+    # Translate single letters to full level names
+    level_map = {
+        "D": "DEBUG",
+        "I": "INFO",
+        "W": "WARNING",
+        "E": "ERROR",
+        "F": "CRITICAL",
+    }
+
+    return level_map.get(env_level, env_level)
+
+
+def custom_timestamper(_: Any, __: Any, event_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """Add ISO UTC timestamp with milliseconds to event_dict['timestamp']."""
+    now = datetime.now()
+    event_dict["timestamp"] = now.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
+    return event_dict
+
+
+class MainConsoleRenderer:
+    """Custom console renderer for [L TIMESTAMP] message with color by level."""
+
+    LEVEL_MAP = {
+        "debug": ("D", "\033[90m"),  # grey
+        "info": ("I", "\033[32m"),  # green
+        "warning": ("W", "\033[33m"),  # yellow
+        "warn": ("W", "\033[33m"),  # yellow
+        "error": ("E", "\033[31m"),  # red
+        "critical": ("F", "\033[41m"),  # red background
+        "fatal": ("F", "\033[41m"),  # alias for critical
+    }
+    RESET = "\033[0m"
+
+    def __init__(self, colors: bool = True):
+        self.colors = colors
+
+    def __call__(
+        self, logger: Any, method_name: str, event_dict: Dict[str, Any]
+    ) -> str:
+        timestamp = event_dict.get("timestamp", "")
+        message = event_dict.get("event", "")
+        level = event_dict.get("level", method_name).lower()
+        letter, color = self.LEVEL_MAP.get(level, ("?", ""))
+        prefix = f"[{letter} {timestamp}]"
+        if self.colors and color:
+            prefix = f"{color}{prefix}{self.RESET}"
+
+        # Build the output with message and key-value pairs
+        output_parts = [prefix]
+
+        # Make the main message bold
+        if self.colors:
+            message = f"\033[1m{message}\033[0m"  # bold
+        output_parts.append(message)
+
+        # Add key-value pairs (excluding internal structlog keys)
+        kv_pairs = []
+        for key, value in event_dict.items():
+            if key not in ["timestamp", "event", "level"]:
+                # Pretty-format complex values (dict/list) as JSON on new lines
+                pretty_value = None
+                if isinstance(value, (dict, list)):
+                    try:
+                        pretty_value = json.dumps(
+                            value, ensure_ascii=False, sort_keys=True, indent=2
+                        )
+                    except Exception:
+                        pretty_value = pformat(value, width=100, compact=False)
+                elif not isinstance(value, (str, int, float, bool, type(None))):
+                    # Fall back to reasonably readable representation for other complex types
+                    pretty_value = pformat(value, width=100, compact=False)
+
+                rendered_value = (
+                    pretty_value if pretty_value is not None else str(value)
+                )
+
+                # If multiline, place value on a new line for readability
+                if "\n" in rendered_value:
+                    if self.colors:
+                        kv_pairs.append(
+                            f"\033[35m{key}\033[0m=\n\033[36m{rendered_value}\033[0m"
+                        )
+                    else:
+                        kv_pairs.append(f"{key}=\n{rendered_value}")
+                else:
+                    if self.colors:
+                        # Format: magenta key + equals + cyan value
+                        kv_pairs.append(
+                            f"\033[35m{key}\033[0m=\033[36m{rendered_value}\033[0m"
+                        )
+                    else:
+                        # No colors for plain output
+                        kv_pairs.append(f"{key}={rendered_value}")
+
+        if kv_pairs:
+            # If any kv is multiline, join with newlines; otherwise keep single line
+            if any("\n" in kv for kv in kv_pairs):
+                kv_text = "\n".join(kv_pairs)
+            else:
+                kv_text = " ".join(kv_pairs)
+            if self.colors:
+                kv_text = f"\033[35m{kv_text}{self.RESET}"  # magenta
+            output_parts.append(kv_text)
+
+        return " ".join(output_parts)
+
+
+def _configure_structlog() -> None:
+    """Configure structlog for both console and file output."""
+    log_dir = _ensure_log_dir()
+    log_file = log_dir / "main.log"
+    json_log_file = log_dir / "main.log.json"
+
+    shared_processors = [
+        structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+        redact_processor,
+        custom_timestamper,
+        structlog.stdlib.add_log_level,
+        structlog.processors.StackInfoRenderer(),
+        structlog.dev.set_exc_info,
+        structlog.processors.format_exc_info,
+        structlog.processors.UnicodeDecoder(),
+    ]
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                # Formatter for colored console output
+                "colored": {
+                    "()": "structlog.stdlib.ProcessorFormatter",
+                    "processors": [
+                        *shared_processors,
+                        MainConsoleRenderer(colors=True),
+                    ],
+                },
+                # Formatter for plain file output
+                "plain": {
+                    "()": "structlog.stdlib.ProcessorFormatter",
+                    "processors": [
+                        *shared_processors,
+                        MainConsoleRenderer(colors=False),
+                    ],
+                },
+                # Formatter for JSON file output
+                "json": {
+                    "()": "structlog.stdlib.ProcessorFormatter",
+                    "processors": [
+                        *shared_processors,
+                        structlog.processors.JSONRenderer(),
+                    ],
+                },
+            },
+            "handlers": {
+                "console": {
+                    "class": "logging.StreamHandler",
+                    "level": _get_env_log_level(),
+                    "formatter": "colored",
+                    "stream": sys.stderr,
+                },
+                "file": {
+                    "class": "logging.handlers.WatchedFileHandler",
+                    "level": "DEBUG",
+                    "filename": log_file,
+                    "formatter": "plain",
+                },
+                "json_file": {
+                    "class": "logging.handlers.WatchedFileHandler",
+                    "level": "DEBUG",
+                    "filename": json_log_file,
+                    "formatter": "json",
+                },
+            },
+            "loggers": {
+                "": {
+                    "handlers": ["console", "file", "json_file"],
+                    "level": "DEBUG",
+                    "propagate": True,
+                },
+            },
+        }
+    )
+
+    structlog.configure(
+        processors=[
+            structlog.stdlib.filter_by_level,
+            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+        ],
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+
+    structlog.get_logger().debug("Logger configured", config=structlog.get_config())
+
+
+# Configure logging on module import
+_configure_structlog()
+
+
+def get_logger(name: str | None = None) -> Any:
+    """Get a configured structlog logger."""
+    return structlog.get_logger(name)
+
+
+# Export the root logger for convenience
+logger = get_logger()

nemo_evaluator_launcher/common/mapping.py

@@ -0,0 +1,295 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import pathlib
+import sys
+from importlib import resources
+from typing import Any, Optional
+
+import requests
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+from nemo_evaluator_launcher.common.logging_utils import logger
+
+# Configuration constants
+# For below, see docs: https://docs.github.com/en/rest/repos/contents
+MAPPING_URL = "https://raw.githubusercontent.com/NVIDIA-NeMo/Eval/main/packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/resources/mapping.toml"
+CACHE_DIR = pathlib.Path.home() / ".nemo-evaluator" / "cache"
+CACHE_FILENAME = "mapping.toml"
+INTERNAL_RESOURCES_PKG = "nemo_evaluator_launcher_internal.resources"
+
+
+def _ensure_cache_dir() -> None:
+    """Ensure the cache directory exists."""
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _get_cache_file() -> pathlib.Path:
+    """Get the cache file path.
+
+    Returns:
+        pathlib.Path: Path to the cache file.
+    """
+    return CACHE_DIR / CACHE_FILENAME
+
+
+def _download_latest_mapping() -> Optional[bytes]:
+    """Download latest mapping from MAPPING_URL and return raw bytes.
+
+    Returns:
+        Optional[bytes]: Downloaded mapping bytes, or None if download fails.
+    """
+    try:
+        response = requests.get(MAPPING_URL, timeout=10)
+        response.raise_for_status()
+
+        # For GitHub raw URLs, the response content is the file content directly
+        mapping_bytes = response.content
+        assert isinstance(mapping_bytes, bytes)
+
+        logger.debug("Successfully downloaded mapping from remote URL")
+        return mapping_bytes
+    except (requests.RequestException, OSError) as e:
+        logger.warning("Failed to download mapping from remote URL", error=str(e))
+        return None
+
+
+def _load_cached_mapping() -> Optional[dict[Any, Any]]:
+    """Load mapping from cache file.
+
+    Returns:
+        Optional[dict]: Loaded mapping data, or None if loading fails.
+    """
+    cache_file = _get_cache_file()
+    if not cache_file.exists():
+        return None
+
+    try:
+        with open(cache_file, "rb") as f:
+            mapping = tomllib.load(f)
+        logger.debug("Loaded mapping from cache")
+        return mapping  # type: ignore[no-any-return]
+    except (OSError, tomllib.TOMLDecodeError) as e:
+        logger.warning("Failed to load mapping from cache", error=str(e))
+        return None
+
+
+def _save_mapping_to_cache(mapping_bytes: bytes) -> None:
+    """Save mapping to cache file.
+
+    Args:
+        mapping_bytes: Mapping data to save.
+    """
+    try:
+        _ensure_cache_dir()
+        cache_file = _get_cache_file()
+
+        # Save the mapping data
+        with open(cache_file, "wb") as f:
+            f.write(mapping_bytes)
+
+    except OSError as e:
+        logger.warning("Failed to save mapping to cache", error=str(e))
+
+
+def _load_packaged_resource(
+    resource_name: str, pkg_name: str = "nemo_evaluator_launcher.resources"
+) -> dict[str, Any]:
+    """Load a resource from the packaged resources.
+
+    Args:
+        resource_name: The name of the resource to load.
+    """
+    try:
+        resource_toml: dict[str, Any] = {}
+        with resources.files(pkg_name).joinpath(resource_name).open("rb") as f:
+            resource_toml = tomllib.load(f)
+        logger.info(
+            "Loaded resource from packaged file", resource=resource_name, pkg=pkg_name
+        )
+        return resource_toml
+    except (OSError, tomllib.TOMLDecodeError) as e:
+        logger.error(
+            "Failed to load from packaged file",
+            resource=resource_name,
+            pkg=pkg_name,
+            error=str(e),
+        )
+        raise RuntimeError(f"Failed to load {resource_name} from packaged file") from e
+
+
+def _process_mapping(mapping_toml: dict) -> dict:
+    """Process the raw mapping TOML into the expected format.
+
+    Args:
+        mapping_toml: Raw mapping TOML data.
+    Returns:
+        dict: Processed mapping in the expected format.
+    """
+    mapping = {}
+    for harness_name, harness_data in mapping_toml.items():
+        assert isinstance(harness_data["tasks"], dict)
+        for endpoint_type, harness_tasks in harness_data["tasks"].items():
+            assert isinstance(harness_tasks, dict)
+            for task_name, task_data in harness_tasks.items():
+                assert isinstance(task_data, dict)
+                key = (harness_name, task_name)
+                if key in mapping:
+                    raise KeyError(
+                        f"(harness,task)-tuple key {repr(key)} already exists in the mapping"
+                    )
+                mapping[key] = {
+                    "task": task_name,
+                    "harness": harness_name,
+                    "container": harness_data["container"],
+                    "endpoint_type": endpoint_type,
+                }
+                for task_data_key in task_data.keys():
+                    if task_data_key in mapping[key]:
+                        raise KeyError(
+                            f"{repr(task_data_key)} is not allowed as key under {repr(key)} in the mapping"
+                        )
+                mapping[key].update(task_data)
+    return mapping
+
+
+def load_tasks_mapping(
+    latest: bool = False,
+    mapping_toml: pathlib.Path | str | None = None,
+) -> dict[tuple[str, str], dict]:
+    """Load tasks mapping.
+
+    The function obeys the following priority rules:
+    1. (Default) If latest==False and mapping_toml is None -> load packaged mapping.
+    2. If latest==True -> fetch MAPPING_URL, save to cache, load it.
+    3. If mapping_toml is not None -> load mapping from this path.
+
+    Returns:
+        dict: Mapping of (harness_name, task_name) to dict holding their configuration.
+
+    """
+    local_mapping: dict = {}
+    if latest:
+        mapping_bytes = _download_latest_mapping()
+        if mapping_bytes:
+            _save_mapping_to_cache(mapping_bytes)
+            local_mapping = _process_mapping(
+                tomllib.loads(mapping_bytes.decode("utf-8"))
+            )
+        else:
+            # Fallback to cached mapping; raise only if cache is missing/invalid
+            cached = _load_cached_mapping()
+            if cached:
+                local_mapping = _process_mapping(cached)
+            else:
+                raise RuntimeError("could not download latest mapping")
+
+    elif mapping_toml is not None:
+        with open(mapping_toml, "rb") as f:
+            local_mapping = _process_mapping(tomllib.load(f))
+    else:
+        local_mapping = _process_mapping(_load_packaged_resource(CACHE_FILENAME))
+
+    # TODO: make more elegant. We consider it ok to avoid a fully-blown plugin system.
+    # Check if nemo_evaluator_launcher_internal is available and load its mapping.toml
+    # CAVEAT: lazy-loading here, not somewhere top level, is important, to ensure
+    # order of package initialization.
+    try:
+        importlib.import_module("nemo_evaluator_launcher_internal")
+        logger.debug("Internal package available, loading internal mapping")
+        internal_mapping = _process_mapping(
+            _load_packaged_resource(CACHE_FILENAME, INTERNAL_RESOURCES_PKG)
+        )
+
+        # Merge internal mapping with local mapping (internal takes precedence)
+        local_mapping.update(internal_mapping)
+        logger.info(
+            "Successfully merged internal mapping", internal_tasks=len(internal_mapping)
+        )
+    except ImportError:
+        logger.debug("Internal package not available, using external mapping only")
+    except Exception as e:
+        logger.warning("Failed to load internal mapping", error=str(e))
+
+    return local_mapping
+
+
+def get_task_from_mapping(query: str, mapping: dict[Any, Any]) -> dict[Any, Any]:
+    """Unambiguously selects one task from the mapping based on the query.
+
+    Args:
+        query: Either `task_name` or `harness_name.task_name`.
+        mapping: The object returned from `load_tasks_mapping` function.
+
+    Returns:
+        dict: Task data.
+
+    """
+    num_dots = query.count(".")
+
+    # if there are no dots in query, treat it like a task name
+    if num_dots == 0:
+        matching_keys = [key for key in mapping.keys() if key[1] == query]
+        # if exactly one task matching the query has been found:
+        if len(matching_keys) == 1:
+            key = matching_keys[0]
+            return mapping[key]  # type: ignore[no-any-return]
+        # if more than one task matching the query has been found:
+        elif len(matching_keys) > 1:
+            matching_queries = [
+                f"{harness_name}.{task_name}"
+                for harness_name, task_name in matching_keys
+            ]
+            raise ValueError(
+                f"there are multiple tasks named {repr(query)} in the mapping,"
+                f" please select one of {repr(matching_queries)}"
+            )
+        # no tasks have been found:
+        else:
+            raise ValueError(f"task {repr(query)} does not exist in the mapping")
+
+    # if there is one dot in query, treat it like "{harness_name}.{task_name}"
+    elif num_dots == 1:
+        harness_name, task_name = query.split(".")
+        matching_keys = [
+            key for key in mapping.keys() if key == (harness_name, task_name)
+        ]
+        # if exactly one task matching the query has been found:
+        if len(matching_keys) == 1:
+            key = matching_keys[0]
+            return mapping[key]  # type: ignore[no-any-return]
+        # if more than one task matching the query has been found:
+        elif len(matching_keys) >= 2:
+            raise ValueError(
+                f"there are multiple matches for {repr(query)} in the mapping,"
+                " which means the mapping is not correct"
+            )
+        # no tasks have been found:
+        else:
+            raise ValueError(
+                f"harness.task {repr(query)} does not exist in the mapping"
+            )
+
+    # invalid query
+    else:
+        raise ValueError(
+            f"invalid query={repr(query)} for task mapping,"
+            " it must contain exactly zero or one occurrence of '.' character"
+        )