# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. 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 hashlib
import os
import subprocess
import tempfile
import time
from typing import Any, Literal, TypeVar
import requests
import yaml
from jinja2 import Environment, StrictUndefined, nodes
from nemo_evaluator.logging import get_logger
__all__ = []
logger = get_logger(__name__)
class MisconfigurationError(Exception):
pass
KeyType = TypeVar("KeyType")
def get_jinja2_environment() -> Environment:
"""Get a configured Jinja2 environment for template operations.
This ensures consistency between template parsing and rendering.
Uses StrictUndefined to match the behavior in api_dataclasses.py.
Returns:
Environment: Configured Jinja2 environment
"""
return Environment(undefined=StrictUndefined)
def deep_update(
mapping: dict[KeyType, Any],
*updating_mappings: dict[KeyType, Any],
skip_nones: bool = False,
) -> dict[KeyType, Any]:
"""Deep update a mapping with other mappings.
If `skip_nones` is True, then the values that are None in the updating mappings are
not updated.
"""
updated_mapping = mapping.copy()
for updating_mapping in updating_mappings:
for k, v in updating_mapping.items():
if (
k in updated_mapping
and isinstance(updated_mapping[k], dict)
and isinstance(v, dict)
):
updated_mapping[k] = deep_update(
updated_mapping[k], v, skip_nones=skip_nones
)
else:
if skip_nones and v is None:
continue
updated_mapping[k] = v
return updated_mapping
def extract_params_from_command(command: str) -> tuple[set[str], set[str]]:
"""Extract all config.params.* parameter names used in a command template.
Args:
command: Jinja2 command template string
Returns:
Tuple of (standard_params, extra_params) where:
- standard_params: Set of param names like {'temperature', 'max_new_tokens'}
- extra_params: Set of params.extra names like {'dummy_score', 'another_param'}
"""
# Use Jinja2's AST parser to extract variable attribute access patterns
# Uses the same environment configuration as template rendering
env = get_jinja2_environment()
ast = env.parse(command)
standard_params = set()
extra_params = set()
def extract_getattr_path(node):
"""Recursively extract the full dotted path from a Getattr node."""
if isinstance(node, nodes.Name):
return node.name
elif isinstance(node, nodes.Getattr):
base = extract_getattr_path(node.node)
if base:
return f"{base}.{node.attr}"
return node.attr
return None
def visit_node(node):
"""Visit all nodes in the AST to find variable references."""
if isinstance(node, nodes.Getattr):
full_path = extract_getattr_path(node)
if full_path:
# Check for config.params.extra.PARAM_NAME pattern
if full_path.startswith("config.params.extra."):
param_name = full_path.replace("config.params.extra.", "")
if param_name: # Only add if there's something after "extra."
# Extract only the first-level key after "extra."
param_name = param_name.split(".")[0]
extra_params.add(param_name)
# Check for config.params.PARAM_NAME pattern (but not just config.params.extra)
elif (
full_path.startswith("config.params.")
and full_path != "config.params.extra"
):
param_name = full_path.replace("config.params.", "")
if param_name != "extra": # Don't add "extra" itself
standard_params.add(param_name)
# Recursively visit child nodes
for child in node.iter_child_nodes():
visit_node(child)
visit_node(ast)
return standard_params, extra_params
def validate_params_in_command(
command: str,
merged_config: dict[KeyType, Any],
) -> None:
"""Validate that all params keys in merged config are used in the command.
Args:
command: The command template from framework.yml
merged_config: The final merged configuration
Raises:
MisconfigurationError: If merged_config contains params keys not used in command
"""
# Extract params keys used in command
command_standard_params, command_extra_params = extract_params_from_command(command)
# Get params from merged config
config_params = merged_config.get("config", {}).get("params", {})
if not config_params:
return # No params to validate
# Check standard params
unused_standard = []
for key, value in config_params.items():
if key == "extra":
continue # Handle extra separately
# Only validate non-None values (None means not set/using default)
if value is not None and key not in command_standard_params:
unused_standard.append(f"config.params.{key}")
# Check params.extra
config_extra = config_params.get("extra", {})
unused_extra = []
for key in config_extra.keys():
if key not in command_extra_params:
unused_extra.append(f"config.params.extra.{key}")
# Raise error if any unused params found
all_unused = unused_standard + unused_extra
if all_unused:
valid_standard = [f"config.params.{p}" for p in sorted(command_standard_params)]
valid_extra = [f"config.params.extra.{p}" for p in sorted(command_extra_params)]
logger.warn(
f"Configuration contains parameter(s) that are not used in the command template: "
f"{', '.join(all_unused)}. "
f"Valid params from command: {valid_standard + valid_extra}. "
f"Remove the unused parameters or update the command template to use them."
)
def dotlist_to_dict(dotlist: list[str]) -> dict:
"""Resolve dot-list style key-value pairs with YAML.
Helper for overriding configuration values using command-line arguments in dot-list style.
"""
dotlist_dict = {}
for override in dotlist:
parts = override.strip().split("=", 1)
if len(parts) == 2:
key = parts[0].strip()
raw_value = parts[1].strip()
# If the value starts with a quote but doesn't end with the same quote,
# it means we have a malformed string. In this case, we'll treat it as a raw string.
if (raw_value.startswith('"') and not raw_value.endswith('"')) or (
raw_value.startswith("'") and not raw_value.endswith("'")
):
value = raw_value
else:
try:
value = yaml.safe_load(raw_value)
except yaml.YAMLError:
# If YAML parsing fails, treat it as a raw string
value = raw_value
keys = key.split(".")
temp = dotlist_dict
for k in keys[:-1]:
temp = temp.setdefault(k, {})
temp[keys[-1]] = value
return dotlist_dict
def run_command(command, cwd=None, verbose=False, propagate_errors=False):
if verbose:
logger.info(f"Running command: {command}")
if cwd:
print(f"Current working directory set to: {cwd}")
with tempfile.TemporaryDirectory() as tmpdirname:
if verbose:
logger.info(f"Temporary directory created at: {tmpdirname}")
file = os.path.join(
tmpdirname, hashlib.sha1(command.encode("utf-8")).hexdigest() + ".sh"
)
if verbose:
logger.info(f"Script file created: {file}")
with open(file, "w") as f:
f.write(command)
f.flush()
if verbose:
logger.info("Command written to script file.")
master, slave = os.openpty()
process = subprocess.Popen(
f"bash {file}",
stdout=slave,
stderr=slave,
stdin=subprocess.PIPE,
cwd=cwd,
shell=True,
executable="/bin/bash",
)
if verbose:
logger.info("Subprocess started.")
os.close(slave)
if propagate_errors:
stderr_output = []
while True:
try:
output = os.read(master, 1024)
if not output:
break
decoded_output = output.decode(errors="ignore")
print(decoded_output, end="", flush=True)
if propagate_errors:
stderr_output.append(decoded_output)
except OSError as e:
if e.errno == 5: # Input/output error is expected at the end of output
break
raise
if verbose:
logger.info("Output reading completed.")
rc = process.wait()
if verbose:
logger.info(f"Subprocess finished with return code: {rc}")
# New error propagation logic
if rc != 0 and propagate_errors:
error_content = (
"".join(stderr_output) if stderr_output else "No error details captured"
)
raise RuntimeError(
f"Evaluation failed! Please consult the logs below:\n{error_content}"
)
return rc
def check_health(
health_url: str, max_retries: int = 600, retry_interval: int = 2
) -> bool:
"""
Check the health of the server.
"""
for _ in range(max_retries):
try:
response = requests.get(health_url)
if response.status_code == 200:
return True
logger.info(f"Server replied with status code: {response.status_code}")
time.sleep(retry_interval)
except requests.exceptions.RequestException:
logger.info("Server is not ready")
time.sleep(retry_interval)
return False
[docs]
def check_endpoint(
endpoint_url: str,
endpoint_type: Literal["completions", "chat"],
model_name: str,
max_retries: int = 600,
retry_interval: int = 2,
) -> bool:
"""Checks if the OpenAI-compatible endpoint is alive by sending a simple prompt.
Args:
endpoint_url (str): Full endpoint URL. For most servers that means either ``/v1/chat/completions`` or ``/completions`` must be provided
endpoint_type (Literal[completions, chat]): indicates if the model is instruction-tuned (chat) or a base model (completions). Used to constuct a proper payload structure.
model_name (str): model name that is linked to payload. Might be required by some endpoint.
max_retries (int, optional): How many attempt before returning false. Defaults to 600.
retry_interval (int, optional): How many seconds to wait between attempts. Defaults to 2.
Raises:
ValueError: if endpoint_type was not one of "completions", "chat"
Returns:
bool: whether the endpoint is alive
"""
payload = {"model": model_name, "max_tokens": 1}
if endpoint_type == "completions":
payload["prompt"] = "hello, my name is"
elif endpoint_type == "chat":
payload["messages"] = [{"role": "user", "content": "hello, what is your name?"}]
else:
raise ValueError(f"Invalid endpoint type: {endpoint_type}")
for _ in range(max_retries):
try:
response = requests.post(endpoint_url, json=payload)
if response.status_code == 200:
return True
logger.info(f"Server replied with status code: {response.status_code}")
time.sleep(retry_interval)
except requests.exceptions.RequestException:
logger.info("Server is not ready")
time.sleep(retry_interval)
return False
