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

nemo_evaluator_launcher/common/mapping.py

@@ -0,0 +1,310 @@
+# 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 base64
+import importlib
+import os
+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.gitlab.com/api/repository_files/
+MAPPING_URL = "TODO: set to github actual one"
+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.
+    """
+    raise NotImplementedError("This logic is still not implemented")
+    try:
+        # Get GitLab token from environment
+        gitlab_token = os.environ.get("GITLAB_TOKEN", "")
+        if not gitlab_token:
+            logger.warning(
+                "GITLAB_TOKEN not set, download may fail due to authentication"
+            )
+        headers = {"PRIVATE-TOKEN": gitlab_token, "Content-Type": "application/json"}
+
+        response = requests.get(MAPPING_URL, headers=headers, timeout=10)
+        response.raise_for_status()
+
+        # GitLab API returns JSON with content in base64
+        response_data = response.json()
+        if "content" in response_data:
+            mapping_bytes = base64.b64decode(response_data["content"])
+        else:
+            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"
+        )

nemo_evaluator_launcher/configs/__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.
+#

nemo_evaluator_launcher/configs/default.yaml

@@ -0,0 +1,28 @@
+# 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.
+#
+defaults:
+  - execution: local
+  - deployment: none
+  - _self_
+
+# NOTE(dfridman): If deployment is used, `target` parameters will be automatically populated.
+target:
+  api_endpoint:
+    url: ???
+    model_id: ???
+    api_key_name: "<YOUR_API_KEY_NAME>" # NOTE: the name of the env var
+
+evaluation: []

nemo_evaluator_launcher/configs/deployment/nim.yaml

@@ -0,0 +1,32 @@
+# 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.
+#
+type: nim
+image: ??? # e.g., nvcr.io/nim/meta/llama-3.1-8b-instruct:1.8.6
+served_model_name: ???
+port: 8000
+
+# NIM containers use default entrypoint - no custom command needed
+# Configuration is done via environment variables in lepton_config
+
+endpoints:
+  chat: /v1/chat/completions
+  completions: /v1/completions
+  health: /health
+# Note: Environment variables should be configured in lepton_config.envs
+# Auto-derived environment variables from deployment config:
+# - SERVED_MODEL_NAME (from served_model_name)
+# - NIM_MODEL_NAME (from served_model_name for NIM)
+# - MODEL_PORT (from port)

nemo_evaluator_launcher/configs/deployment/none.yaml

@@ -0,0 +1,16 @@
+# 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.
+#
+type: none

nemo_evaluator_launcher/configs/deployment/sglang.yaml

@@ -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.
+#
+type: sglang
+image: lmsysorg/sglang:latest
+checkpoint_path: ???
+served_model_name: ???
+port: 8000
+tensor_parallel_size: 8
+data_parallel_size: 1
+extra_args: ""
+env_vars: {} # {name: value} dict
+
+endpoints:
+  chat: /v1/chat/completions
+  completions: /v1/completions
+  health: /health
+
+command: python3 -m sglang.launch_server
+  --model-path ${oc.select:deployment.hf_model_handle,/checkpoint}
+  --host 0.0.0.0
+  --port ${deployment.port}
+  --served-model-name ${deployment.served_model_name}
+  --tp ${deployment.tensor_parallel_size}
+  --dp ${deployment.data_parallel_size}
+  ${deployment.extra_args}

nemo_evaluator_launcher/configs/deployment/vllm.yaml

@@ -0,0 +1,41 @@
+# 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.
+#
+type: vllm
+image: vllm/vllm-openai:latest
+checkpoint_path: ???
+served_model_name: ???
+port: 8000
+tensor_parallel_size: 8
+pipeline_parallel_size: 1
+data_parallel_size: 1
+extra_args: ""
+env_vars: {} # {name: value} dict
+
+endpoints:
+  chat: /v1/chat/completions
+  completions: /v1/completions
+  health: /health
+
+command: vllm serve ${oc.select:deployment.hf_model_handle,/checkpoint}
+  --tensor-parallel-size=${deployment.tensor_parallel_size}
+  --pipeline-parallel-size=${deployment.pipeline_parallel_size}
+  --data-parallel-size=${deployment.data_parallel_size}
+  --port ${deployment.port}
+  --trust-remote-code
+  --served-model-name ${deployment.served_model_name}
+  --enforce-eager
+  --gpu-memory-utilization 0.95
+  ${deployment.extra_args}

nemo_evaluator_launcher/configs/execution/lepton/default.yaml

@@ -0,0 +1,92 @@
+# 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.
+#
+# =============================================================================
+# LEPTON EXECUTION CONFIGURATION - Default Environment
+# Following the same pattern as SLURM internal configs
+# =============================================================================
+
+type: lepton
+output_dir: ???
+
+# Environment variables passed to evaluation containers
+env_var_names: []
+
+# Evaluation task settings for Lepton batch execution
+evaluation_tasks:
+  # Resource requirements for evaluation tasks
+  resource_shape: "cpu.small"  # Can be overridden to gpu.small, gpu.1xh200, etc
+
+  # Timeout for individual evaluation tasks (in seconds)
+  timeout: 3600  # 1 hour per evaluation task
+
+  # Whether to use shared storage for results
+  use_shared_storage: true
+
+# Lepton platform infrastructure settings for this environment
+lepton_platform:
+  # Deployment configuration (for model endpoints)
+  deployment:
+    # Node group for endpoint deployments
+    node_group: "nv-int-multiteam-nebius-h200-01-mjgbgffo"
+
+    # Endpoint readiness timeout (in seconds)
+    endpoint_readiness_timeout: 1200  # 20 minutes
+
+    # Platform defaults that get merged into all deployments
+    platform_defaults:
+      # Health checks
+      health:
+        readiness: {}
+        liveness: {}
+
+      # Logging
+      log:
+        enable_collection: true
+
+      # Monitoring and routing
+      metrics: {}
+      routing_policy: {}
+
+      # Queue configuration for resource prioritization
+      queue_config:
+        priority_class: "mid-4000"
+
+      # Security and infrastructure
+      enable_rdma: false
+      user_security_context: {}
+
+      # NGC registry access
+      image_pull_secrets:
+        - "ngc-general-access-and-private-registry-in-sa-nvex-org"
+
+  # Task execution configuration (for evaluation tasks)
+  tasks:
+    # Node group for evaluation tasks
+    node_group: "nv-int-multiteam-nebius-h200-01"
+
+    # Default environment variables for all tasks
+    env_vars: {}
+
+    # Storage mounts for task execution
+    mounts:
+      # Main workspace mount
+      - from: "node-nfs:lepton-shared-fs"
+        path: "/shared/nemo-evaluator-launcher-workspace"
+        mount_path: "/workspace"
+
+    # Image pull secrets for task containers
+    image_pull_secrets:
+      - "lepton-nvidia"

nemo_evaluator_launcher/configs/execution/local.yaml

@@ -0,0 +1,17 @@
+# 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.
+#
+type: local
+output_dir: ???

nemo_evaluator_launcher/configs/execution/slurm/default.yaml

@@ -0,0 +1,33 @@
+# 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.
+#
+# Each slurm cluster has its own flavour, below we provide some defaults that might meet one's needs.
+hostname: ???
+username: ${oc.env:USER}
+account: ???
+partition: batch
+num_nodes: 1
+ntasks_per_node: 1
+gres: gpu:8
+walltime: 01:00:00
+subproject: nemo-evaluator-launcher
+output_dir: ???
+env_vars:
+  deployment: {}
+  evaluation: {}
+mounts:
+  deployment: {}
+  evaluation: {}
+  mount_home: true

nemo_evaluator_launcher/executors/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+#
+# NOTE(dfridman): importing all executors will register them in the registry
+
+from nemo_evaluator_launcher.executors.lepton.executor import LeptonExecutor
+from nemo_evaluator_launcher.executors.local.executor import LocalExecutor
+from nemo_evaluator_launcher.executors.slurm.executor import SlurmExecutor
+
+__all__ = ["LeptonExecutor", "LocalExecutor", "SlurmExecutor"]