nemo-evaluator-launcher 0.1.0rc2__py3-none-any.whl

nemo_evaluator_launcher/executors/local/executor.py

@@ -0,0 +1,491 @@
+# 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.
+#
+"""Local executor implementation for nemo-evaluator-launcher.
+
+Handles running evaluation jobs locally using shell scripts and Docker containers.
+"""
+
+import copy
+import os
+import pathlib
+import platform
+import shlex
+import subprocess
+import time
+from typing import List, Optional
+
+import jinja2
+import yaml
+from omegaconf import DictConfig, OmegaConf
+
+from nemo_evaluator_launcher.common.execdb import (
+    ExecutionDB,
+    JobData,
+    generate_invocation_id,
+    generate_job_id,
+)
+from nemo_evaluator_launcher.common.helpers import (
+    get_eval_factory_command,
+    get_eval_factory_dataset_size_from_run_config,
+    get_timestamp_string,
+)
+from nemo_evaluator_launcher.common.mapping import (
+    get_task_from_mapping,
+    load_tasks_mapping,
+)
+from nemo_evaluator_launcher.executors.base import (
+    BaseExecutor,
+    ExecutionState,
+    ExecutionStatus,
+)
+from nemo_evaluator_launcher.executors.registry import register_executor
+
+
+@register_executor("local")
+class LocalExecutor(BaseExecutor):
+    @classmethod
+    def execute_eval(cls, cfg: DictConfig, dry_run: bool = False) -> str:
+        """Run evaluation jobs locally using the provided configuration.
+
+        Args:
+            cfg: The configuration object for the evaluation run.
+            dry_run: If True, prepare scripts and save them without execution.
+
+        Returns:
+            str: The invocation ID for the evaluation run.
+
+        Raises:
+            NotImplementedError: If deployment is not 'none'.
+            RuntimeError: If the run script fails.
+        """
+        if cfg.deployment.type != "none":
+            raise NotImplementedError(
+                f"type {cfg.deployment.type} is not implemented -- add deployment support"
+            )
+
+        # Generate invocation ID for this evaluation run
+        invocation_id = generate_invocation_id()
+
+        output_dir = pathlib.Path(cfg.execution.output_dir).absolute() / (
+            get_timestamp_string(include_microseconds=False) + "-" + invocation_id
+        )
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        tasks_mapping = load_tasks_mapping()
+        evaluation_tasks = []
+        job_ids = []
+
+        eval_template = jinja2.Template(
+            open(pathlib.Path(__file__).parent / "run.template.sh", "r").read()
+        )
+
+        execution_mode = cfg.execution.get("mode", "parallel")
+        if execution_mode == "parallel":
+            is_execution_mode_sequential = False
+        elif execution_mode == "sequential":
+            is_execution_mode_sequential = True
+        else:
+            raise ValueError(
+                "unknown execution mode: {}. Choose one of {}".format(
+                    repr(execution_mode), ["parallel", "sequential"]
+                )
+            )
+
+        for idx, task in enumerate(cfg.evaluation.tasks):
+            task_definition = get_task_from_mapping(task.name, tasks_mapping)
+
+            # Create job ID as <invocation_id>.<n>
+            job_id = generate_job_id(invocation_id, idx)
+            job_ids.append(job_id)
+            container_name = f"{task.name}-{get_timestamp_string()}"
+
+            # collect all env vars
+            env_vars = copy.deepcopy(dict(cfg.evaluation.get("env_vars", {})))
+            env_vars.update(task.get("env_vars", {}))
+            if cfg.target.api_endpoint.api_key_name:
+                assert "API_KEY" not in env_vars
+                env_vars["API_KEY"] = cfg.target.api_endpoint.api_key_name
+
+            # check if the environment variables are set
+            for env_var in env_vars.values():
+                if os.getenv(env_var) is None:
+                    raise ValueError(
+                        f"Trying to pass an unset environment variable {env_var}."
+                    )
+
+            # check if required env vars are defined:
+            for required_env_var in task_definition.get("required_env_vars", []):
+                if required_env_var not in env_vars.keys():
+                    raise ValueError(
+                        f"{task.name} task requires environment variable {required_env_var}."
+                        " Specify it in the task subconfig in the 'env_vars' dict as the following"
+                        f" pair {required_env_var}: YOUR_ENV_VAR_NAME"
+                    )
+
+            # format env_vars for a template
+            env_vars = [
+                f"{env_var_dst}=${env_var_src}"
+                for env_var_dst, env_var_src in env_vars.items()
+            ]
+
+            eval_image = task_definition["container"]
+            if "container" in task:
+                eval_image = task["container"]
+
+            task_output_dir = output_dir / task.name
+            task_output_dir.mkdir(parents=True, exist_ok=True)
+            evaluation_task = {
+                "name": task.name,
+                "job_id": job_id,
+                "eval_image": eval_image,
+                "container_name": container_name,
+                "env_vars": env_vars,
+                "output_dir": task_output_dir,
+                "eval_factory_command": get_eval_factory_command(
+                    cfg, task, task_definition
+                ),
+            }
+            evaluation_tasks.append(evaluation_task)
+
+            # Check if auto-export is enabled by presence of destination(s)
+            auto_export_config = cfg.execution.get("auto_export", {})
+            auto_export_destinations = auto_export_config.get("destinations", [])
+
+            run_sh_content = (
+                eval_template.render(
+                    evaluation_tasks=[evaluation_task],
+                    auto_export_destinations=auto_export_destinations,
+                ).rstrip("\n")
+                + "\n"
+            )
+
+            (task_output_dir / "run.sh").write_text(run_sh_content)
+
+        run_all_sequentially_sh_content = (
+            eval_template.render(
+                evaluation_tasks=evaluation_tasks,
+                auto_export_destinations=auto_export_destinations,
+            ).rstrip("\n")
+            + "\n"
+        )
+        (output_dir / "run_all.sequential.sh").write_text(
+            run_all_sequentially_sh_content
+        )
+
+        # Save launched jobs metadata
+        db = ExecutionDB()
+        for job_id, task, evaluation_task in zip(
+            job_ids, cfg.evaluation.tasks, evaluation_tasks
+        ):
+            db.write_job(
+                job=JobData(
+                    invocation_id=invocation_id,
+                    job_id=job_id,
+                    timestamp=time.time(),
+                    executor="local",
+                    data={
+                        "output_dir": str(evaluation_task["output_dir"]),
+                        "container": evaluation_task["container_name"],
+                        "eval_image": evaluation_task["eval_image"],
+                    },
+                    config=OmegaConf.to_object(cfg),
+                )
+            )
+
+        if dry_run:
+            print("\n\n=============================================\n\n")
+            print(f"DRY RUN: Scripts prepared and saved to {output_dir}")
+            if is_execution_mode_sequential:
+                print(
+                    "\n\n =========== Main script | run_all.sequential.sh ===================== \n\n"
+                )
+                with open(output_dir / "run_all.sequential.sh", "r") as f:
+                    print(f.read())
+            else:
+                for idx, task in enumerate(cfg.evaluation.tasks):
+                    task_output_dir = output_dir / task.name
+                    print(
+                        f"\n\n =========== Task script | {task.name}/run.sh ===================== \n\n"
+                    )
+                    with open(task_output_dir / "run.sh", "r") as f:
+                        print(f.read())
+            print("\nTo execute, run without --dry-run")
+            return invocation_id
+
+        # Launch bash scripts with Popen for non-blocking execution.
+        # To ensure subprocess continues after python exits:
+        # - on Unix-like systems, to fully detach the subprocess
+        #   so it does not die when Python exits, pass start_new_session=True;
+        # - on Widnows use creationflags=subprocess.CREATE_NEW_PROCESS_GROUP flag.
+        os_name = platform.system()
+        if is_execution_mode_sequential:
+            if os_name == "Windows":
+                subprocess.Popen(
+                    shlex.split("bash run_all.sequential.sh"),
+                    cwd=output_dir,
+                    creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+                )
+            else:
+                subprocess.Popen(
+                    shlex.split("bash run_all.sequential.sh"),
+                    cwd=output_dir,
+                    start_new_session=True,
+                )
+        else:
+            for task in cfg.evaluation.tasks:
+                if os_name == "Windows":
+                    subprocess.Popen(
+                        shlex.split("bash run.sh"),
+                        cwd=output_dir / task.name,
+                        creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+                    )
+                else:
+                    subprocess.Popen(
+                        shlex.split("bash run.sh"),
+                        cwd=output_dir / task.name,
+                        start_new_session=True,
+                    )
+
+        print("\nCommands for real-time monitoring:")
+        for job_id, evaluation_task in zip(job_ids, evaluation_tasks):
+            log_file = evaluation_task["output_dir"] / "logs" / "stdout.log"
+            print(f"  tail -f {log_file}")
+
+        print("\nFollow all logs for this invocation:")
+        print(f"  tail -f {output_dir}/*/logs/stdout.log")
+
+        return invocation_id
+
+    @staticmethod
+    def get_status(id: str) -> List[ExecutionStatus]:
+        """Get the status of a specific job or all jobs in an invocation group.
+
+        Args:
+            id: Unique job identifier or invocation identifier.
+
+        Returns:
+            List containing the execution status for the job(s).
+        """
+        db = ExecutionDB()
+
+        # If id looks like an invocation_id (8 hex digits, no dot), get all jobs for it
+        if len(id) == 8 and "." not in id:
+            jobs = db.get_jobs(id)
+            statuses: List[ExecutionStatus] = []
+            for job_id, _ in jobs.items():
+                statuses.extend(LocalExecutor.get_status(job_id))
+            return statuses
+
+        # Otherwise, treat as job_id
+        job_data = db.get_job(id)
+        if job_data is None:
+            return []
+        if job_data.executor != "local":
+            return []
+
+        output_dir = pathlib.Path(job_data.data.get("output_dir", ""))
+        if not output_dir.exists():
+            return [ExecutionStatus(id=id, state=ExecutionState.PENDING)]
+
+        artifacts_dir = output_dir / "artifacts"
+        progress = _get_progress(artifacts_dir)
+
+        logs_dir = output_dir / "logs"
+        if not logs_dir.exists():
+            return [
+                ExecutionStatus(
+                    id=id,
+                    state=ExecutionState.PENDING,
+                    progress=dict(progress=progress),
+                )
+            ]
+
+        # Check if job was killed
+        if job_data.data.get("killed", False):
+            return [
+                ExecutionStatus(
+                    id=id, state=ExecutionState.KILLED, progress=dict(progress=progress)
+                )
+            ]
+
+        stage_files = {
+            "pre_start": logs_dir / "stage.pre-start",
+            "running": logs_dir / "stage.running",
+            "exit": logs_dir / "stage.exit",
+        }
+
+        if stage_files["exit"].exists():
+            try:
+                content = stage_files["exit"].read_text().strip()
+                if " " in content:
+                    timestamp, exit_code_str = content.rsplit(" ", 1)
+                    exit_code = int(exit_code_str)
+                    if exit_code == 0:
+                        return [
+                            ExecutionStatus(
+                                id=id,
+                                state=ExecutionState.SUCCESS,
+                                progress=dict(progress=progress),
+                            )
+                        ]
+                    else:
+                        return [
+                            ExecutionStatus(
+                                id=id,
+                                state=ExecutionState.FAILED,
+                                progress=dict(progress=progress),
+                            )
+                        ]
+                else:
+                    return [
+                        ExecutionStatus(
+                            id=id,
+                            state=ExecutionState.FAILED,
+                            progress=dict(progress=progress),
+                        )
+                    ]
+            except (ValueError, OSError):
+                return [
+                    ExecutionStatus(
+                        id=id,
+                        state=ExecutionState.FAILED,
+                        progress=dict(progress=progress),
+                    )
+                ]
+        elif stage_files["running"].exists():
+            return [
+                ExecutionStatus(
+                    id=id,
+                    state=ExecutionState.RUNNING,
+                    progress=dict(progress=progress),
+                )
+            ]
+        elif stage_files["pre_start"].exists():
+            return [
+                ExecutionStatus(
+                    id=id,
+                    state=ExecutionState.PENDING,
+                    progress=dict(progress=progress),
+                )
+            ]
+
+        return [
+            ExecutionStatus(
+                id=id, state=ExecutionState.PENDING, progress=dict(progress=progress)
+            )
+        ]
+
+    @staticmethod
+    def kill_job(job_id: str) -> None:
+        """Kill a local job by stopping its Docker container and related processes.
+
+        Args:
+            job_id: The job ID to kill.
+
+        Raises:
+            ValueError: If job is not found or invalid.
+            RuntimeError: If Docker container cannot be stopped.
+        """
+        db = ExecutionDB()
+        job_data = db.get_job(job_id)
+
+        if job_data is None:
+            raise ValueError(f"Job {job_id} not found")
+
+        if job_data.executor != "local":
+            raise ValueError(
+                f"Job {job_id} is not a local job (executor: {job_data.executor})"
+            )
+
+        # Get container name from database
+        container_name = job_data.data.get("container")
+        if not container_name:
+            raise ValueError(f"No container name found for job {job_id}")
+
+        killed_something = False
+
+        # First, try to stop the Docker container if it's running
+        result = subprocess.run(
+            shlex.split(f"docker stop {container_name}"),
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode == 0:
+            killed_something = True
+        # Don't raise error if container doesn't exist (might be still pulling)
+
+        # Find and kill Docker processes for this container
+        result = subprocess.run(
+            shlex.split(f"pkill -f 'docker run.*{container_name}'"),
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode == 0:
+            killed_something = True
+
+        # Mark job as killed in database if we killed something
+        if killed_something:
+            job_data.data["killed"] = True
+            db.write_job(job_data)
+        else:
+            raise RuntimeError(
+                f"Could not find or kill job {job_id} (container: {container_name})"
+            )
+
+
+def _get_progress(artifacts_dir: pathlib.Path) -> Optional[float]:
+    """Get the progress of a local job.
+
+    Args:
+        artifacts_dir: The directory containing the evaluation artifacts.
+
+    Returns:
+        The progress of the job as a float between 0 and 1.
+    """
+    progress_filepath = artifacts_dir / "progress"
+    if not progress_filepath.exists():
+        return None
+    progress_str = progress_filepath.read_text().strip()
+    try:
+        processed_samples = int(progress_str)
+    except ValueError:
+        return None
+
+    dataset_size = _get_dataset_size(artifacts_dir)
+    if dataset_size is not None:
+        progress = processed_samples / dataset_size
+    else:
+        # NOTE(dfridman): if we don't know the dataset size, report the number of processed samples
+        progress = processed_samples
+    return progress
+
+
+def _get_dataset_size(artifacts_dir: pathlib.Path) -> Optional[int]:
+    """Get the dataset size for a benchmark.
+
+    Args:
+        artifacts_dir: The directory containing the evaluation artifacts.
+
+    Returns:
+        The dataset size for the benchmark.
+    """
+    run_config = artifacts_dir / "run_config.yml"
+    if not run_config.exists():
+        return None
+    run_config = yaml.safe_load(run_config.read_text())
+    return get_eval_factory_dataset_size_from_run_config(run_config)

nemo_evaluator_launcher/executors/local/run.template.sh

@@ -0,0 +1,88 @@
+# 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.
+
+
+# check if docker exists
+command -v docker >/dev/null 2>&1 || { echo 'docker not found'; exit 1; }
+
+{% for task in evaluation_tasks %}
+# {{ task.job_id }} {{ task.name }}
+
+task_dir="{{ task.output_dir }}"
+artifacts_dir="$task_dir/artifacts"
+logs_dir="$task_dir/logs"
+
+mkdir -m 777 -p "$task_dir"
+mkdir -m 777 -p "$artifacts_dir"
+mkdir -m 777 -p "$logs_dir"
+
+# Create pre-start stage file
+echo "$(date -u +%Y-%m-%dT%H:%M:%SZ)" > "$logs_dir/stage.pre-start"
+
+# Docker run with eval factory command
+(
+    echo "$(date -u +%Y-%m-%dT%H:%M:%SZ)" > "$logs_dir/stage.running"
+    docker run --rm --shm-size=100g \
+      --name {{ task.container_name }} \
+      --volume "$artifacts_dir":/results \
+      {% for env_var in task.env_vars -%}
+      -e {{ env_var }} \
+      {% endfor -%}
+      {{ task.eval_image }} \
+      bash -c '
+        {{ task.eval_factory_command }} ;
+        exit_code=$?
+        chmod 777 -R /results;
+        if [ "$exit_code" -ne 0 ]; then
+            echo "The evaluation container failed with exit code $exit_code" >&2;
+            exit "$exit_code";
+        fi;
+        echo "Container completed successfully" >&2;
+        exit 0;
+      ' > "$logs_dir/stdout.log" 2>&1
+    exit_code=$?
+    echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) $exit_code" > "$logs_dir/stage.exit"
+) >> "$logs_dir/stdout.log" 2>&1
+
+
+{% if auto_export_destinations %}
+# Monitor job completion and auto-export
+(
+    # Give it a moment to ensure file is fully written
+    sleep 1
+
+    exit_code=$(tail -1 "$logs_dir/stage.exit" | cut -d' ' -f2)
+    if [ "$exit_code" = "0" ]; then
+        # Log auto-export activity to task logs
+        echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Job {{ task.job_id }} completed successfully. Starting auto-export..." >> "$logs_dir/stdout.log"
+
+        {% for dest in auto_export_destinations %}
+        echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Exporting job {{ task.job_id }} to {{ dest }}..." >> "$logs_dir/stdout.log"
+        nemo-evaluator-launcher export {{ task.job_id }} --dest {{ dest }} >> "$logs_dir/stdout.log" 2>&1
+        if [ $? -eq 0 ]; then
+            echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Export to {{ dest }} completed successfully" >> "$logs_dir/stdout.log"
+        else
+            echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Export to {{ dest }} failed" >> "$logs_dir/stdout.log"
+        fi
+        {% endfor %}
+
+        echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Auto-export completed for job {{ task.job_id }}" >> "$logs_dir/stdout.log"
+    else
+        echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) Job {{ task.job_id }} failed with exit code $exit_code. Skipping auto-export." >> "$logs_dir/stdout.log"
+    fi
+)
+
+{% endif %}
+{% endfor %}

nemo_evaluator_launcher/executors/registry.py

@@ -0,0 +1,38 @@
+# 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.
+#
+from typing import Callable, Type
+
+from nemo_evaluator_launcher.executors.base import BaseExecutor
+
+_EXECUTOR_REGISTRY: dict[str, Type[BaseExecutor]] = {}
+
+
+def register_executor(
+    executor_name: str,
+) -> Callable[[Type[BaseExecutor]], Type[BaseExecutor]]:
+    def wrapper(executor_cls: Type[BaseExecutor]) -> Type[BaseExecutor]:
+        _EXECUTOR_REGISTRY[executor_name] = executor_cls
+        return executor_cls
+
+    return wrapper
+
+
+def get_executor(executor_name: str) -> Type[BaseExecutor]:
+    if executor_name not in _EXECUTOR_REGISTRY:
+        raise ValueError(
+            f"Executor {executor_name} not found. Available executors: {list(_EXECUTOR_REGISTRY.keys())}"
+        )
+    return _EXECUTOR_REGISTRY[executor_name]

nemo_evaluator_launcher/executors/slurm/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+#