nemo-evaluator-launcher 0.1.19__py3-none-any.whl → 0.1.56__py3-none-any.whl

nemo_evaluator_launcher/api/functional.py

@@ -18,8 +18,9 @@
 This module provides the main functional entry points for running evaluations, querying job status, and listing available tasks. These functions are intended to be used by CLI commands and external integrations.
 """
 
+import copy
 from pathlib import Path
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import yaml
 from omegaconf import DictConfig, OmegaConf
@@ -35,7 +36,7 @@ def get_tasks_list() -> list[list[Any]]:
     """Get a list of available tasks from the mapping.
 
     Returns:
-        list[list[Any]]: Each sublist contains task name, endpoint type, harness, and container.
+        list[list[Any]]: Each sublist contains task name, endpoint type, harness, container, arch, description, and type.
     """
     mapping = load_tasks_mapping()
     data = [
@@ -44,6 +45,9 @@ def get_tasks_list() -> list[list[Any]]:
             task_data.get("endpoint_type"),
             task_data.get("harness"),
             task_data.get("container"),
+            task_data.get("arch", ""),
+            task_data.get("description", ""),
+            task_data.get("type", ""),
         ]
         for task_data in mapping.values()
     ]
@@ -75,12 +79,54 @@ def _validate_no_missing_values(cfg: Any, path: str = "") -> None:
             _validate_no_missing_values(value, current_path)
 
 
-def run_eval(cfg: RunConfig, dry_run: bool = False) -> Optional[str]:
+def filter_tasks(cfg: RunConfig, task_names: list[str]) -> RunConfig:
+    """Filter evaluation tasks to only include specified task names.
+
+    Args:
+        cfg: The configuration object for the evaluation run.
+        task_names: List of task names to include (e.g., ["ifeval", "gsm8k"]).
+
+    Returns:
+        RunConfig: A new configuration with filtered tasks (input is not mutated).
+
+    Raises:
+        ValueError: If any requested task is not found in config or no tasks defined.
+    """
+    if not task_names:
+        return cfg
+
+    if not hasattr(cfg.evaluation, "tasks") or not cfg.evaluation.tasks:
+        raise ValueError("No tasks defined in config. Cannot filter tasks.")
+
+    requested_tasks = set(task_names)
+    original_tasks = cfg.evaluation.tasks
+    filtered_tasks = [task for task in original_tasks if task.name in requested_tasks]
+
+    # Fail if ANY requested tasks are not found
+    found_names = {task.name for task in filtered_tasks}
+    not_found = requested_tasks - found_names
+    if not_found:
+        available = [task.name for task in original_tasks]
+        raise ValueError(
+            f"Requested task(s) not found in config: {sorted(not_found)}. "
+            f"Available tasks: {available}"
+        )
+
+    # Create a deep copy to preserve input immutability
+    result = copy.deepcopy(cfg)
+    result.evaluation.tasks = filtered_tasks
+    return result
+
+
+def run_eval(
+    cfg: RunConfig, dry_run: bool = False, tasks: Optional[list[str]] = None
+) -> Optional[str]:
     """Run evaluation with specified config and overrides.
 
     Args:
         cfg: The configuration object for the evaluation run.
         dry_run: If True, do not run the evaluation, just prepare scripts and save them.
+        tasks: Optional list of task names to run. If provided, only these tasks will be executed.
 
     Returns:
         Optional[str]: The invocation ID for the evaluation run.
@@ -89,6 +135,10 @@ def run_eval(cfg: RunConfig, dry_run: bool = False) -> Optional[str]:
         ValueError: If configuration validation fails or MISSING values are found.
         RuntimeError: If the executor fails to start the evaluation.
     """
+    # Filter tasks if specified
+    if tasks:
+        cfg = filter_tasks(cfg, tasks)
+
     # Validate that no MISSING values exist in the configuration
     _validate_no_missing_values(cfg)
 
@@ -116,6 +166,7 @@ def get_status(ids_or_prefixes: list[str]) -> list[dict[str, Any]]:
     db = ExecutionDB()
     results: List[dict[str, Any]] = []
 
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
     for id_or_prefix in ids_or_prefixes:
         # If id looks like an invocation_id (no dot), get all jobs for it
         if "." not in id_or_prefix:
@@ -259,6 +310,108 @@ def get_status(ids_or_prefixes: list[str]) -> list[dict[str, Any]]:
     return results
 
 
+def stream_logs(
+    ids_or_prefixes: Union[str, list[str]],
+) -> Iterator[Tuple[str, str, str]]:
+    """Stream logs from jobs or invocations by their IDs or invocation IDs.
+
+    Args:
+        ids_or_prefixes: Single ID/prefix or list of job IDs or invocation IDs to stream logs from.
+                         Short prefixes are allowed, we would try to match the full ones from
+                         prefixes if no collisions are present.
+
+    Yields:
+        Tuple[str, str, str]: Tuples of (job_id, task_name, log_line) for each log line.
+            Empty lines are yielded as empty strings.
+
+    Raises:
+        ValueError: If the executor doesn't support log streaming.
+    """
+    db = ExecutionDB()
+
+    # Normalize to list for consistent processing
+    if isinstance(ids_or_prefixes, str):
+        ids_or_prefixes = [ids_or_prefixes]
+
+    # Collect all jobs from all IDs, grouped by executor
+    executor_to_jobs: Dict[str, Dict[str, JobData]] = {}
+    executor_to_invocations: Dict[str, list[str]] = {}
+
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
+    for id_or_prefix in ids_or_prefixes:
+        # Determine if this is a job ID or invocation ID
+        if "." in id_or_prefix:
+            # This is a job ID
+            job_data = db.get_job(id_or_prefix)
+            if job_data is None:
+                continue
+
+            executor = job_data.executor
+            if executor not in executor_to_jobs:
+                executor_to_jobs[executor] = {}
+            executor_to_jobs[executor][id_or_prefix] = job_data
+        else:
+            # This is an invocation ID
+            jobs = db.get_jobs(id_or_prefix)
+            if not jobs:
+                continue
+
+            # Get the executor class from the first job
+            first_job_data = next(iter(jobs.values()))
+            executor = first_job_data.executor
+            if executor not in executor_to_invocations:
+                executor_to_invocations[executor] = []
+            executor_to_invocations[executor].append(id_or_prefix)
+
+    # Stream logs from each executor simultaneously
+    # For each executor, collect all job IDs and stream them together
+    for executor, jobs_dict in executor_to_jobs.items():
+        try:
+            executor_cls = get_executor(executor)
+        except ValueError:
+            continue
+
+        # For local executor with multiple jobs, pass list to stream simultaneously
+        # For other executors or single jobs, pass individual job IDs
+        if executor == "local" and len(jobs_dict) > 1:
+            # Pass all job IDs as a list to stream simultaneously
+            try:
+                yield from executor_cls.stream_logs(
+                    list(jobs_dict.keys()), executor_name=executor
+                )
+            except NotImplementedError:
+                raise ValueError(
+                    f"Log streaming is not yet implemented for executor '{executor}'"
+                )
+        else:
+            # Single job or non-local executor
+            for job_id in jobs_dict.keys():
+                try:
+                    yield from executor_cls.stream_logs(job_id, executor_name=executor)
+                except NotImplementedError:
+                    raise ValueError(
+                        f"Log streaming is not yet implemented for executor '{executor}'"
+                    )
+
+    # Stream logs from invocation IDs
+    for executor, invocation_ids in executor_to_invocations.items():
+        try:
+            executor_cls = get_executor(executor)
+        except ValueError:
+            continue
+
+        # Stream each invocation (each invocation already handles multiple jobs internally)
+        for invocation_id in invocation_ids:
+            try:
+                yield from executor_cls.stream_logs(
+                    invocation_id, executor_name=executor
+                )
+            except NotImplementedError:
+                raise ValueError(
+                    f"Log streaming is not yet implemented for executor '{executor}'"
+                )
+
+
 def list_all_invocations_summary() -> list[dict[str, Any]]:
     """Return a concise per-invocation summary from the exec DB.
 
@@ -378,6 +531,7 @@ def kill_job_or_invocation(id: str) -> list[dict[str, Any]]:
                 "data": {"error": f"Unexpected error: {str(e)}"},
             }
 
+    # TODO(agronskiy): refactor the `.`-checking job in all the functions.
     # Determine if this is a job ID or invocation ID
     if "." in id:
         # This is a job ID - kill single job
@@ -413,7 +567,7 @@ def kill_job_or_invocation(id: str) -> list[dict[str, Any]]:
 
 
 def export_results(
-    invocation_ids: Union[str, List[str]],
+    invocation_ids: Union[str, list[str]],
     dest: str = "local",
     config: dict[Any, Any] | None = None,
 ) -> dict:
@@ -442,7 +596,7 @@ def export_results(
             if "." in single_id:  # job_id
                 # Try reading config from artifacts working dir (auto-export on remote node)
                 cfg_file = None
-                for name in ("run_config.yml", "config.yml"):
+                for name in ("config.yml", "run_config.yml"):
                     p = Path(name)
                     if p.exists():
                         cfg_file = p

nemo_evaluator_launcher/cli/logs.py

@@ -0,0 +1,102 @@
+# 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.
+#
+"""Logs command for streaming logs from evaluation jobs."""
+
+import sys
+from dataclasses import dataclass
+from typing import Callable, Dict
+
+from simple_parsing import field
+
+import nemo_evaluator_launcher.common.printing_utils as pu
+from nemo_evaluator_launcher.api.functional import stream_logs
+from nemo_evaluator_launcher.common.execdb import ExecutionDB
+from nemo_evaluator_launcher.common.logging_utils import logger
+
+
+@dataclass
+class Cmd:
+    """Logs command configuration."""
+
+    ids: list[str] = field(
+        default_factory=list,
+        positional=True,
+        help="Invocation IDs or job IDs (e.g., '15b9f667' or '15b9f667.0'). Multiple IDs can be provided.",
+    )
+
+    def execute(self) -> None:
+        """Execute the logs command to stream logs from jobs."""
+        if not self.ids:
+            logger.error("At least one ID is required")
+            sys.exit(1)
+
+        db = ExecutionDB()
+
+        # Validate all IDs exist
+        all_job_ids = []
+        for id_or_prefix in self.ids:
+            if "." in id_or_prefix:
+                # This is a job ID - get single job
+                job_data = db.get_job(id_or_prefix)
+                if job_data is None:
+                    logger.error(f"Job {id_or_prefix} not found")
+                    sys.exit(1)
+                all_job_ids.append(id_or_prefix)
+            else:
+                # This is an invocation ID - get all jobs
+                jobs = db.get_jobs(id_or_prefix)
+                if not jobs:
+                    logger.error(f"Invocation {id_or_prefix} not found")
+                    sys.exit(1)
+                all_job_ids.extend(jobs.keys())
+
+        # Build color mapping for job IDs
+        colors = [pu.red, pu.green, pu.yellow, pu.magenta, pu.cyan]
+        job_colors: Dict[str, Callable[[str], str]] = {}
+        color_index = 0
+
+        for job_id in all_job_ids:
+            job_colors[job_id] = colors[color_index % len(colors)]
+            color_index += 1
+
+        # Stream logs from executor
+        try:
+            log_stream = stream_logs(self.ids)
+            for job_id, task_name, log_line in log_stream:
+                # Extract short prefix: first 6 chars of invocation ID + job number
+                if "." in job_id:
+                    inv_id, job_num = job_id.split(".", 1)
+                    short_prefix = f"{inv_id[:6]}.{job_num}"
+                else:
+                    short_prefix = job_id[:6]
+                prefix = f"{short_prefix}:"
+                color_func = job_colors.get(job_id, pu.grey)
+                if log_line:
+                    print(f"{color_func(prefix)} {log_line}")
+                else:
+                    # Print empty lines without prefix
+                    print()
+
+        except ValueError:
+            # Handle case where executor doesn't support streaming
+            # Warning already logged by BaseExecutor.stream_logs
+            pass
+        except KeyboardInterrupt:
+            # Clean exit on Ctrl+C
+            pass
+        except Exception as e:
+            logger.error(f"Error streaming logs: {e}")
+            sys.exit(1)

nemo_evaluator_launcher/cli/ls_task.py

@@ -0,0 +1,280 @@
+# 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.
+#
+"""CLI command for listing task details."""
+
+import json
+from dataclasses import dataclass
+
+import yaml
+from simple_parsing import field
+
+from nemo_evaluator_launcher.common.container_metadata import (
+    TaskIntermediateRepresentation,
+    load_tasks_from_tasks_file,
+)
+from nemo_evaluator_launcher.common.logging_utils import logger
+from nemo_evaluator_launcher.common.mapping import load_tasks_mapping
+from nemo_evaluator_launcher.common.printing_utils import (
+    bold,
+    cyan,
+    magenta,
+    yellow,
+)
+
+
+@dataclass
+class Cmd:
+    """List task command configuration."""
+
+    task_identifier: str = field(
+        default="",
+        positional=True,
+        help="Task identifier in format '[harness.]task_name'. If empty, shows all tasks.",
+    )
+    json: bool = field(
+        default=False,
+        action="store_true",
+        help="Print output as JSON instead of formatted text",
+    )
+    tasks_file: str = field(
+        default="",
+        help="Path to all_tasks_irs.yaml file (default: auto-detect)",
+    )
+    from_container: str = field(
+        default="",
+        help="Load tasks from container image (e.g., nvcr.io/nvidia/eval-factory/simple-evals:25.10). "
+        "If provided, extracts framework.yml from container and loads tasks on-the-fly instead of using all_tasks_irs.yaml",
+    )
+
+    def execute(self) -> None:
+        """Execute the ls task command."""
+        import pathlib
+
+        # Initialize tasks_path to None - it will be set when loading from file
+        tasks_path = None
+
+        # If --from is provided, load tasks from container
+        if self.from_container:
+            from nemo_evaluator_launcher.common.container_metadata import (
+                load_tasks_from_container,
+            )
+
+            try:
+                tasks = load_tasks_from_container(self.from_container)
+            except ValueError as e:
+                print(f"Error: {e}")
+                return
+            except Exception as e:
+                logger.error(
+                    "Failed to load tasks from container",
+                    container=self.from_container,
+                    error=str(e),
+                    exc_info=True,
+                )
+                return
+
+            if not tasks:
+                logger.error(
+                    "No tasks found in container",
+                    container=self.from_container,
+                )
+                return
+
+            logger.debug(
+                "Loaded tasks from container",
+                container=self.from_container,
+                num_tasks=len(tasks),
+                containers=set(task.container for task in tasks),
+            )
+            mapping_verified = True  # Tasks from container are always verified
+        else:
+            # Default behavior: load from all_tasks_irs.yaml
+            if self.tasks_file:
+                tasks_path = pathlib.Path(self.tasks_file)
+                if not tasks_path.exists():
+                    logger.error("Tasks file not found", path=str(tasks_path))
+                    return
+
+            # Load tasks
+            try:
+                tasks, mapping_verified = load_tasks_from_tasks_file(tasks_path)
+            except Exception as e:
+                print(f"Error loading tasks: {e}")
+                import traceback
+
+                traceback.print_exc()
+                logger.error("Failed to load tasks", error=str(e), exc_info=True)
+                return
+
+            # Display warning if mapping is not verified
+            if not mapping_verified:
+                print(
+                    yellow(
+                        "⚠ Warning: Tasks are from unverified mapping (mapping.toml checksum mismatch)"
+                    )
+                )
+                print(
+                    yellow(
+                        "  Consider regenerating all_tasks_irs.yaml if mapping.toml has changed"
+                    )
+                )
+                print()
+
+            # Override containers from mapping.toml (which has the latest containers)
+            # This ensures ls task shows the same containers as ls tasks
+            # Only do this when NOT using --from (when loading from all_tasks_irs.yaml)
+            try:
+                mapping = load_tasks_mapping()
+                # Create a lookup: (normalized_harness, normalized_task_name) -> container
+                # Use case-insensitive keys for matching
+                container_lookup = {}
+                for (harness, task_name), task_data in mapping.items():
+                    container = task_data.get("container")
+                    if container:
+                        # Normalize harness name for lookup (frameworks.yaml uses hyphens)
+                        normalized_harness = harness.replace("_", "-").lower()
+                        normalized_task = task_name.lower()
+                        container_lookup[(normalized_harness, normalized_task)] = (
+                            container
+                        )
+
+                # Update task containers from mapping.toml
+                for task in tasks:
+                    # Defensive checks: ensure task has required attributes
+                    if not hasattr(task, "harness") or not task.harness:
+                        logger.warning(
+                            "Task missing harness attribute, skipping container override",
+                            task_name=getattr(task, "name", "unknown"),
+                        )
+                        continue
+                    if not hasattr(task, "name") or not task.name:
+                        logger.warning(
+                            "Task missing name attribute, skipping container override",
+                            harness=getattr(task, "harness", "unknown"),
+                        )
+                        continue
+
+                    # Normalize both harness and task name for case-insensitive lookup
+                    normalized_harness = task.harness.lower()
+                    normalized_task = task.name.lower()
+                    lookup_key = (normalized_harness, normalized_task)
+                    if lookup_key in container_lookup:
+                        task.container = container_lookup[lookup_key]
+            except Exception as e:
+                logger.debug(
+                    "Failed to override containers from mapping.toml",
+                    error=str(e),
+                )
+                # Continue with containers from all_tasks_irs.yaml if mapping load fails
+
+        if not tasks:
+            print("No tasks found.")
+            if tasks_path:
+                print(f"  Tasks file: {tasks_path}")
+            else:
+                print(
+                    "  Note: Make sure all_tasks_irs.yaml exists and contains valid task definitions."
+                )
+            return
+
+        # Parse task identifier
+        harness_filter = None
+        task_filter = None
+        if self.task_identifier:
+            if "." in self.task_identifier:
+                parts = self.task_identifier.split(".", 1)
+                harness_filter = parts[0]
+                task_filter = parts[1]
+            else:
+                task_filter = self.task_identifier
+
+        # Filter tasks
+        filtered_tasks = []
+        for task in tasks:
+            if harness_filter and task.harness.lower() != harness_filter.lower():
+                continue
+            if task_filter and task.name.lower() != task_filter.lower():
+                continue
+            filtered_tasks.append(task)
+
+        if not filtered_tasks:
+            print(f"No tasks found matching: {self.task_identifier}")
+            if self.task_identifier:
+                # Show available tasks for debugging
+                print("\nAvailable tasks (showing first 10):")
+                for i, task in enumerate(tasks[:10]):
+                    print(f"  - {task.harness}.{task.name}")
+                if len(tasks) > 10:
+                    print(f"  ... and {len(tasks) - 10} more")
+            return
+
+        # Display tasks
+        if self.json:
+            self._print_json(filtered_tasks)
+        else:
+            self._print_formatted(filtered_tasks, mapping_verified)
+
+    def _print_json(self, tasks: list[TaskIntermediateRepresentation]) -> None:
+        """Print tasks as JSON."""
+        tasks_dict = [task.to_dict() for task in tasks]
+        print(json.dumps({"tasks": tasks_dict}, indent=2))
+
+    def _print_formatted(
+        self, tasks: list[TaskIntermediateRepresentation], mapping_verified: bool = True
+    ) -> None:
+        """Print tasks in formatted text with colorized output."""
+        for i, task in enumerate(tasks):
+            if i > 0:
+                print()  # Spacing between tasks
+                print(bold("=" * 80))
+
+            # Task name - bold and magenta key, cyan value (matching logging utils)
+            print(f"{bold(magenta('Task:'))} {bold(cyan(str(task.name)))}")
+
+            # Description - magenta key, cyan value
+            if task.description:
+                print(f"{magenta('Description:')} {cyan(str(task.description))}")
+
+            # Harness - magenta key, cyan value
+            print(f"{magenta('Harness:')} {cyan(str(task.harness))}")
+
+            # Container - magenta key, cyan value
+            print(f"{magenta('Container:')} {cyan(str(task.container))}")
+
+            # Container Digest - magenta key, cyan value
+            if task.container_digest:
+                print(
+                    f"{magenta('Container Digest:')} {cyan(str(task.container_digest))}"
+                )
+
+            # Print defaults as YAML
+            if task.defaults:
+                print(f"\n{bold(magenta('Defaults:'))}")
+                defaults_yaml = yaml.dump(
+                    task.defaults, default_flow_style=False, sort_keys=False
+                )
+                # Indent defaults - use cyan for YAML content (FDF values)
+                for line in defaults_yaml.split("\n"):
+                    if line.strip():
+                        print(f"  {cyan(line)}")
+                    else:
+                        print()
+
+            print(bold("-" * 80))
+
+        # Total count - bold
+        task_word = "task" if len(tasks) == 1 else "tasks"
+        print(f"\n{bold(f'Total: {len(tasks)} {task_word}')}")